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 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>,
}

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