cog - 专门用于将机器学习模型打包进生产级容器的开源工具

Cog 项目介绍

Cog是一款开源工具，专为将机器学习模型打包成标准化、可用于生产的容器而设计。利用Cog，用户可以轻松地将已打包的模型部署到自己的基础设施，或是使用Replicate来进行部署。

项目亮点

📦 免除Docker文件的困扰：自己编写Dockerfile或许让人困惑，而Cog通过一个简单的配置文件定义环境，并自动生成包含最佳实践的Docker镜像。包括Nvidia基础镜像、依赖项的高效缓存、特定Python版本的安装及合理预设的环境变量等。
🤬️ 杜绝CUDA配置的麻烦：Cog可以识别CUDA/cuDNN/PyTorch/Tensorflow/Python的兼容组合，并自动为用户配置。
✅ 用标准Python定义模型的输入和输出：Cog会生成OpenAPI模式，并用Pydantic验证输入和输出。
🎁 自动HTTP预测服务器：借助FastAPI，根据模型的类型动态生成RESTful HTTP API。
🥞 自动队列工作者：对于长时间运行的深度学习模型或批处理，队列是最佳的架构。Cog模型自带此功能。目前支持Redis，未来将支持更多。
☁️ 云存储支持：可直接读取和写入Amazon S3与Google Cloud Storage（即将推出）。
🚀 生产准备就绪：可部署到任何支持Docker镜像的环境，包括用户自己的基础设施或Replicate。

工作原理

用户可以通过cog.yaml文件定义模型运行的Docker环境，例如：

build:
  gpu: true
  system_packages:
    - "libgl1-mesa-glx"
    - "libglib2.0-0"
  python_version: "3.12"
  python_packages:
    - "torch==2.3"
predict: "predict.py:Predictor"

使用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)

通过命令行运行预测：

$ cog predict -i image=@input.jpg
--> 正在构建Docker镜像...
--> 运行预测...
--> 输出保存至output.jpg

或者，构建Docker镜像以供部署：

$ cog build -t my-colorization-model
--> 正在构建Docker镜像...
--> 构建成功 my-colorization-model:latest

$ docker run -d -p 5000:5000 --gpus all my-colorization-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"}}'