Cog：机器学习容器化工具

Cog 是一个开源工具，让你能够以标准、生产就绪的容器形式打包机器学习模型。

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

亮点

• 📦 免去 Docker 配置之痛。 自行编写 Dockerfile 可能是一大难题。使用 Cog，你只需通过简单的配置文件定义环境，它就会生成包含所有最佳实践的 Docker 镜像：Nvidia 基础镜像、高效依赖缓存、安装指定 Python 版本、合理的环境变量默认值等。

• 🤬️ 告别 CUDA 地狱。 Cog 了解哪些 CUDA/cuDNN/PyTorch/TensorFlow/Python 组合是兼容的，并会为你正确配置好一切。

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

• 🎁 自动 HTTP 推理服务器： 根据模型类型，动态生成基于高性能 Rust/Axum 服务器的 RESTful HTTP API。

• 🚀 生产就绪。 将你的模型部署到任何支持 Docker 镜像的环境，无论是自有基础设施还是 Replicate。

工作原理

通过 cog.yaml 定义模型运行的 Docker 环境：

build:
  gpu: true
  system_packages:
    - "libgl1"
    - "libglib2.0-0"
  python_version: "3.13"
  python_requirements: requirements.txt
run: "run.py:Runner"

通过 run.py 定义模型的运行方式：

from cog import BaseRunner, Input, Path
import torch

class Runner(BaseRunner):
    def setup(self):
        """将模型加载到内存中，以便高效执行多次推理"""
        self.model = torch.load("./weights.pth")

    # 模型接受输入的参数和类型
    def run(self,
          image: Path = Input(description="灰度输入图像")
    ) -> Path:
        """运行模型"""
        processed_image = preprocess(image)
        output = self.model(processed_image)
        return postprocess(output)

上述示例中，我们接受图像路径作为输入，经过模型处理后返回转换后图像的路径。

现在，你可以运行模型：

$ cog run -i image=@input.jpg
--> 正在构建 Docker 镜像...
--> 正在运行...
--> 输出已写入 output.jpg

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

$ cog build -t my-classification-model
--> 正在构建 Docker 镜像...
--> 已构建 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 是解决方案的一部分，但要让一切运行起来非常复杂：Dockerfile、前后处理、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 及使用 WSL 2 的 Windows 11
• Docker。Cog 使用 Docker 为你的模型创建容器，需先安装 Docker。如果安装的是 Docker Engine 而非 Docker Desktop，还需安装 Buildx。

安装

根据你的平台选择相应的安装指南。

brew install replicate/tap/cog

# bash、zsh 及其他 shell
sh <(curl -fsSL https://cog.run/install.sh)

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

# 使用 wget 下载并在单独命令中运行
wget -qO- https://cog.run/install.sh
sh ./install.sh

sudo curl -o /usr/local/bin/cog -L "https://github.com/replicate/cog/releases/latest/download/cog_$(uname -s)_$(uname -m | sed 's/aarch64/arm64/')"
sudo chmod +x /usr/local/bin/cog
sudo xattr -d com.apple.quarantine /usr/local/bin/cog 2>/dev/null || true

sudo xattr -d com.apple.quarantine /usr/local/bin/cog

# bash、zsh 及其他 shell
sh <(curl -fsSL https://cog.run/install.sh)

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

# 使用 wget 下载并在单独命令中运行
wget -qO- https://cog.run/install.sh
sh ./install.sh

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

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

升级

如果使用 macOS 且已通过 Homebrew 安装，请运行：

brew upgrade replicate/tap/cog

否则，重复执行安装时使用的命令即可升级到最新版本。

开发

有关设置开发环境和从源码构建的详细信息，请参阅 CONTRIBUTING.md。

下一步

• 从示例模型开始
• 使用你自己的模型
• 在 notebook 中使用 Cog
• 在 Windows 11 上使用 Cog
• 浏览本仓库中的示例模型
• 使用 Cog 部署模型
• cog.yaml 参考文档，了解如何定义模型环境
• 运行接口参考，了解 Runner 接口的工作原理
• 训练接口参考，了解如何为模型添加微调 API
• HTTP API 参考，了解模型提供的 HTTP API 的使用方法

需要帮助？

加入我们的 Discord #cog 频道。

贡献者 ✨

感谢以下杰出人士（表情符号说明）：

[贡献者列表]

本项目遵循 all-contributors 规范。欢迎各类形式的贡献！