Cog：机器学习容器

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

你可以将打包好的模型部署到你自己的基础设施上，或者部署到 Replicate。

亮点

• 📦 无需痛苦的 Docker 容器。 编写自己的 Dockerfile 可能是一个令人困惑的过程。使用 Cog，你只需通过一个简单的配置文件定义你的环境，它就会生成一个包含所有最佳实践的 Docker 镜像：Nvidia 基础镜像、依赖项的高效缓存、安装特定 Python 版本、合理的环境变量默认值等等。

• 🤬️ 告别 CUDA 地狱。 Cog 知道哪些 CUDA/cuDNN/PyTorch/Tensorflow/Python 组合是兼容的，并会为你正确设置一切。

• ✅ 使用标准 Python 定义模型的输入和输出。 然后，Cog 会生成一个 OpenAPI schema (模式) 并验证输入和输出。

• 🎁 自动 HTTP 预测服务器：利用高性能的 Rust/Axum 服务器，根据模型的类型动态生成 RESTful HTTP API。

• 🚀 准备就绪，可用于生产环境。 在任何运行 Docker 镜像的地方部署你的模型。你自己的基础设施，或 Replicate。

工作原理

使用 cog.yaml 定义模型运行的 Docker 环境：

build:
  gpu: true
  system_packages:
    - "libgl1-mesa-glx"
    - "libglib2.0-0"
  python_version: "3.13"
  python_requirements: requirements.txt
predict: "predict.py:Predictor"

使用 predict.py 定义如何在模型上运行预测：

from cog import BasePredictor, Input, Path
import torch

class Predictor(BasePredictor):
    def setup(self):
        """Load the model into memory to make running multiple predictions efficient"""
        self.model = torch.load("./weights.pth")

    # The arguments and types the model takes as input
    def predict(self,
          image: Path = Input(description="Grayscale input image")
    ) -> Path:
        """Run a single prediction on the model"""
        processed_image = preprocess(image)
        output = self.model(processed_image)
        return postprocess(output)

在上面，我们接受图像路径作为输入，并在将其通过模型处理后返回转换后图像的路径。

现在，你可以在此模型上运行预测：

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

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

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

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

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

或者，通过 serve 命令结合构建和运行：

$ cog serve -p 8080

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

我们为什么要构建这个？

研究人员很难将机器学习模型投入生产。

解决方案的一部分是 Docker，但让它工作起来非常复杂：Dockerfiles、预处理/后处理、Flask 服务器、CUDA 版本等。大多数情况下，研究人员不得不与工程师坐下来才能将该东西部署成功。

Andreas 和 Ben 创建了 Cog。Andreas 曾在 Spotify 工作，在那里他构建了使用 Docker 构建和部署 ML (机器学习) 模型的工具。Ben 曾在 Docker 工作，在那里他创建了 Docker Compose。

我们意识到，除了 Spotify，其他公司也在使用 Docker 来构建和部署机器学习模型。Uber 和其他公司已构建了类似的系统。因此，我们正在制作一个开源版本，以便其他人也能这样做。

如果你有兴趣使用它或想与我们合作，请联系我们。我们在 Discord 上，或者发邮件到 team@replicate.com。

前置条件

• macOS、Linux 或 Windows 11。Cog 支持 macOS、Linux 和 Windows 11（配合 WSL 2 (Windows Subsystem for Linux 2)）。
• Docker。Cog 使用 Docker 为你的模型创建容器。你需要先 安装 Docker 才能运行 Cog。如果你安装的是 Docker Engine 而不是 Docker Desktop，你还需要 安装 Buildx。

安装

如果你使用的是 macOS，可以使用 Homebrew 安装 Cog：

brew install replicate/tap/cog

你也可以使用我们的 安装脚本 下载并安装最新版本：

# bash, zsh, and other shells
sh <(curl -fsSL https://cog.run/install.sh)

# fish shell
sh (curl -fsSL https://cog.run/install.sh | psub)

# download with wget and run in a separate command
wget -qO- https://cog.run/install.sh
sh ./install.sh

你也可以直接在终端中运行以下命令，从 GitHub 手动安装 Cog 的最新版本：

sudo curl -o /usr/local/bin/cog -L "https://github.com/replicate/cog/releases/latest/download/cog_$(uname -s)_$(uname -m)"
sudo chmod +x /usr/local/bin/cog

或者如果你是在 Docker 中：

RUN sh -c "INSTALL_DIR=\"/usr/local/bin\" SUDO=\"\" $(curl -fsSL https://cog.run/install.sh)"

升级

如果你使用的是 macOS 并且之前使用 Homebrew 安装了 Cog，请运行以下内容：

brew upgrade replicate/tap/cog

否则，你可以通过运行与安装时相同的命令升级到最新版本。

开发

有关如何设置开发环境和从源代码构建的说明，请参阅 CONTRIBUTING.md。

下一步

• 使用示例模型入门
• 使用自己的模型入门
• 在 Notebooks 中使用 Cog
• 在 Windows 11 上使用 Cog
• 查看一些使用 Cog 的示例
• 使用 Cog 部署模型
• cog.yaml 参考文档 了解如何定义模型的运行环境
• 预测接口参考 了解 Predictor 接口的工作原理
• 训练接口参考 了解如何为模型添加微调 API
• HTTP API 参考 了解如何使用模型提供的 HTTP API

需要帮助？

加入我们的 Discord #cog 频道。

贡献者 ✨

感谢这些杰出的人 (表情符号说明)：

Ben Firshman 💻 📖	Andreas Jansson 💻 📖 🚧	Zeke Sikelianos 💻 📖 🔧	Rory Byrne 💻 📖 ⚠️	Michael Floering 💻 📖 🤔	Ben Evans 📖	shashank agarwal 💻 📖
VictorXLR 💻 📖 ⚠️	hung anna 🐛	Brian Whitman 🐛	JimothyJohn 🐛	ericguizzo 🐛	Dominic Baggott 💻 ⚠️	Dashiell Stander 🐛 💻 ⚠️
Shuwei Liang 🐛 💬	Eric Allam 🤔	Iván Perdomo 🐛	Charles Frye 📖	Luan Pham 🐛 📖	TommyDew 💻	Jesse Andrews 💻 📖 ⚠️
Nick Stenning 💻 📖 🎨 🚇 ⚠️	Justin Merrell 📖	Rurik Ylä-Onnenvuori 🐛	Youka 🐛	Clay Mullis 📖	Mattt 💻 📖 🚇	Eng Zer Jun ⚠️
BB 💻	williamluer 📖	Simon Eskildsen 💻	F 🐛 💻	Philip Potter 🐛 💻	Joanne Chen 📖	technillogue 💻
Aron Carroll 📖 💻 🤔	Bohdan Mykhailenko 📖 🐛	Daniel Radu 📖 🐛	Itay Etelis 💻	Gennaro Schiano 📖	André Knörig 📖	Dan Fairs 💻

本项目遵循 all-contributors（所有贡献者）规范。欢迎任何形式的贡献！

Cog 快速上手指南

Cog 是一个开源工具，允许你将机器学习模型打包为标准的生产就绪容器。你可以将打包好的模型部署到自有基础设施或 Replicate。

环境准备

• 操作系统：macOS、Linux 或 Windows 11（Windows 下需配合 WSL 2）。
• 前置依赖：Docker。Cog 使用 Docker 创建模型容器。
  • 如果你安装的是 Docker Engine 而非 Docker Desktop，还需要安装 Buildx。

安装步骤

根据你的操作系统选择以下一种方式安装：

macOS (Homebrew)

brew install replicate/tap/cog

其他系统 (通用脚本)

# bash, zsh 等
sh <(curl -fsSL https://cog.run/install.sh)

# fish shell
sh (curl -fsSL https://cog.run/install.sh | psub)

手动安装 (GitHub Release)

sudo curl -o /usr/local/bin/cog -L "https://github.com/replicate/cog/releases/latest/download/cog_$(uname -s)_$(uname -m)"
sudo chmod +x /usr/local/bin/cog

基本使用

1. 定义环境 (cog.yaml)

在项目根目录创建 cog.yaml 文件，定义运行环境和依赖：

build:
  gpu: true
  system_packages:
    - "libgl1-mesa-glx"
    - "libglib2.0-0"
  python_version: "3.13"
  python_requirements: requirements.txt
predict: "predict.py:Predictor"

2. 定义预测逻辑 (predict.py)

创建 predict.py 文件，定义模型的输入输出及推理逻辑：

from cog import BasePredictor, Input, Path
import torch

class Predictor(BasePredictor):
    def setup(self):
        """加载模型到内存以提高多次预测效率"""
        self.model = torch.load("./weights.pth")

    def predict(self,
          image: Path = Input(description="灰度输入图像")
    ) -> Path:
        """运行单次预测"""
        processed_image = preprocess(image)
        output = self.model(processed_image)
        return postprocess(output)

3. 运行与构建

本地直接预测

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

构建 Docker 镜像

$ cog build -t my-classification-model
--> Built my-classification-model:latest

启动 HTTP 服务

$ cog serve -p 8080

随后即可通过 HTTP API 调用模型：

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