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

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

其输出应与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]
cpu = 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