Cog[-rust]: 机器学习的容器

Cog是一个开源工具，允许您将Rust机器学习模型打包成标准、生产就绪的容器。

它的输出应该可以与Replicate自己的Cog（用于Python模型）互换。

亮点

• 📦 无需痛苦地编写Docker容器。 编写自己的 Dockerfile 可能是一个令人困惑的过程。使用Cog，您可以在Cargo.toml中定义您的环境，并且它会生成一个包含所有最佳实践的Docker镜像：Nvidia基础镜像、依赖项的高效缓存、最小镜像大小、合理的环境变量默认值等。

• 🤬️ 不再有CUDA地狱。 Cog知道哪些CUDA/cuDNN/tch/tensorflow组合是兼容的，并将为您正确设置一切。

• ✅ 在Rust中定义您的模型的输入和输出。 然后，Cog生成一个OpenAPI模式，并使用JSONSchema验证输入和输出。

• 🎁 自动HTTP预测服务器：使用 axum，您的模型类型将用于动态生成RESTful HTTP API。

• ☁️ 云存储。 可以直接将文件读取和写入到Amazon S3和Google Cloud Storage。（即将推出。）

• 🚀 适用于生产。 在任何运行Docker镜像的地方部署您的模型。您自己的基础设施，或者 Replicate。

它是如何工作的

轻松在您的Cargo.toml中定义您的环境。Cog推断其余部分

[package]
name = "ml-model"

[package.metadata.cog]
gpu = true # optional, defaults to false
image = "docker-image-name" # optional, defaults to `cog-[package.name]`

在您的main.rs中定义如何在您的模型上运行预测

use anyhow::Result;
use async_trait::async_trait;
use cog_rust::Cog;
use schemars::JsonSchema;
use std::collections::HashMap;
use tch::{
  nn::{ModuleT, VarStore},
  vision::{imagenet, resnet::resnet50},
  Device,
};

#[derive(serde::Deserialize, schemars::JsonSchema)]
struct ModelRequest {
  /// Image to classify
  image: cog_rust::Path,
}

struct ResnetModel {
  model: Box<dyn ModuleT + Send>,
}

#[async_trait]
impl Cog for ResnetModel {
  type Request = ModelRequest;
  type Response = HashMap<String, f64>;

  async fn setup() -> Result<Self> {
    let mut vs = VarStore::new(Device::Cpu);
    vs.load("weights/model.safetensors")?;
    let model = Box::new(resnet50(&vs.root(), imagenet::CLASS_COUNT));

    Ok(Self { model })
  }

  fn predict(&self, input: Self::Request) -> Result<Self::Response> {
    let image = imagenet::load_image_and_resize224(&input.image)?;
    let output = self
      .model
      .forward_t(&image.unsqueeze(0), false)
      .softmax(-1, tch::Kind::Float);

    Ok(imagenet::top(&output, 5)
      .into_iter()
      .map(|(prob, class)| (class, 100.0 * prob))
      .collect())
  }
}

cog_rust::start!(ResnetModel);

现在，您可以在该模型上运行预测

$ cargo cog predict -i @input.jpg
--> Building Docker image...
--> Running Prediction...
--> Output written to output.jpg

或者，构建一个用于部署的Docker镜像

$ cargo cog build -t my-colorization-model
--> Building Docker image...
--> Built my-colorization-model:latest

$ docker run -d -p 5000:5000 --gpus all my-colorization-model

$ curl https://127.0.0.1:5000/predictions -X POST \
    -H 'Content-Type: application/json' \
    -d '{"input": {"image": "https://.../input.jpg"}}'

为什么我要构建这个？

Replicate 团队成功地建立了从 Python 笔记本到 Docker 镜像再到 API 端点最简单的方式。

然而，使用 Python 作为基础层也带来了一些挑战，比如巨大的镜像大小或者模型请求时的额外延迟。

随着非 Python 机器学习生态系统的逐渐繁荣（例如，参见 whisper.cpp 和 llama.cpp），cog-rust 将提供在相同的接口上暴露的额外性能，这些接口用户和工具已经熟悉。

先决条件

• macOS、Linux 或 Windows。Cog 在 Rust 可运行的所有地方都可以工作。
• Docker。Cog 使用 Docker 为您的模型创建一个容器。在您运行 Cog 之前，您需要 安装 Docker。

安装

您可以使用 Cargo 安装 Cog

cargo install cargo-cog

用法

$ cargo cog --help
A cargo subcommand to build, run and publish machine learning containers

Usage: cargo cog [OPTIONS] [COMMAND]

Commands:
  login    Log in to Replicate's Docker registry
  build    Build the model in the current directory into a Docker image
  push     Build and push model in current directory to a Docker registry
  predict  Run a prediction
  help     Print this message or the help of the given subcommand(s)

Options:
  -h, --help     Print help
  -V, --version  Print version