在docker容器中部署AI模型

（本文由腾讯混元模型和ChatGPT辅助编写）

在本篇博客中，我们将探讨如何在Docker环境中部署一个基于PyTorch和Transformers库的AI模型，用于图像描述生成。我们将使用一个简单的Flask应用作为后端服务，并通过Docker容器化该应用，以便于部署和管理。

项目概述

本项目实现了一个图像描述生成功能，用户可以通过POST请求上传图片，服务器将返回对应的图像描述。项目代码仓库为duyixian1234/itt。

关键技术栈

Python: 主要编程语言
Flask: 轻量级Web框架
PyTorch & Transformers: 深度学习框架和预训练模型库
Docker: 容器化技术，用于部署和管理应用

推理代码

from flask import Flask, request
from PIL import Image
from transformers import BlipForConditionalGeneration, BlipProcessor

processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
model = BlipForConditionalGeneration.from_pretrained(
    "Salesforce/blip-image-captioning-base"
)
print("Model loaded")

app = Flask(__name__)


@app.post("/recognize")
def recognize():
    file = request.files["file"]
    raw_image = Image.open(file.stream).convert("RGB")
    inputs = processor(raw_image, return_tensors="pt")
    out = model.generate(**inputs, max_new_tokens=256)
    return processor.decode(out[0], skip_special_tokens=True)

这段代码是一个使用Flask框架搭建的Web应用，它实现了一个图像识别功能，基于BLIP（Bootstrapping Language-Image Pre-training）模型进行图片描述生成。具体作用如下：

Flask应用初始化：
- 使用Flask创建了一个Web应用，设置了一个POST请求接口/recognize，用于接收上传的图像文件。
加载模型：
- 通过transformers库加载了预训练的BLIP模型和相应的处理器BlipProcessor，这两个组件用于对图像进行处理并生成描述。加载的模型是Salesforce/blip-image-captioning-base。
/recognize接口：
- 当客户端发送带有图片的POST请求到/recognize接口时：
  - 通过request.files["file"]获取上传的文件。
  - 使用PIL库的Image.open()方法将上传的图片转换为RGB格式。
  - 使用BlipProcessor将图像转换为模型需要的输入格式。
  - 使用BlipForConditionalGeneration模型生成图像的描述文本，最多生成256个新的标记。
  - 最后返回生成的描述，去除特殊标记。
功能概述：
- 客户端上传一张图片，服务器利用BLIP模型自动生成该图片的描述，返回给客户端。

Dockerfile详解

FROM python:3.12-slim-bookworm AS base
COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/

FROM base AS build
WORKDIR /app

# Install dependencies
COPY pyproject.toml uv.lock /app/
RUN uv sync --frozen

# Load model
RUN uv run python -c 'from transformers import BlipForConditionalGeneration, BlipProcessor;BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base");model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")'

# Copy source code
COPY main.py .

FROM build AS deploy
# Run the server
EXPOSE 8000
CMD ["uv", "run", "gunicorn","main:app","-k","gevent"]

我们的Dockerfile分为三个阶段：基础镜像构建、依赖安装与模型加载、最终部署。

基础镜像构建

我们使用python:3.12-slim-bookworm作为基础镜像，并从ghcr.io/astral-sh/uv:latest复制uv工具链，以便于后续的依赖管理和构建。

依赖安装与模型加载

在此阶段，我们首先复制pyproject.toml和uv.lock文件，并使用uv工具同步依赖。然后，我们在Dockerfile中提前加载模型，这样可以减少后续构建时间，因为模型文件较大，下载和加载过程可能较为耗时。

最终部署

在最终部署阶段，我们暴露8000端口，并使用gunicorn作为WSGI服务器启动Flask应用。这里我们使用gevent工作模式，以提高并发处理能力。

PyTorch版本选择

[project]
name = "itt"
version = "0.1.0"
description = "Add your description here"
readme = "README.md"
requires-python = ">=3.12"
dependencies = [
    "flask>=3.1.0",
    "gevent>=24.11.1",
    "gunicorn>=23.0.0",
    "torch>=2.5.1",
    "torchvision>=0.20.1",
    "transformers>=4.46.3",
]

[tool.uv.sources]
torch = { index = "pytorch" }
torchvision = { index = "pytorch" }

[[tool.uv.index]]
name = "pytorch"
url = "https://download.pytorch.org/whl/cpu"
explicit = true

[[tool.uv.index]]
name = "tencent"
url = "https://mirrors.tencent.com/pypi/simple"

在安装PyTorch时，需要根据实际需求选择CPU版或CUDA版。CPU版适用于没有GPU或不需要GPU加速的场景，而CUDA版则适用于具备NVIDIA GPU且需要加速计算的场景。在pyproject.toml文件中，我们通过指定不同的索引源来区分版本。

性能与消耗

在标准的GitHub Codespaces环境中，本镜像运行时大约消耗1G内存。处理一张图片的时间约为6秒，这个时间包括了图像的上传、预处理、模型推理和结果返回等步骤。

结语

通过本文，我们展示了如何在Docker环境中部署一个AI模型，并提供了详细的Dockerfile配置和相关技术栈说明。希望这篇博客能为你在AI模型的部署道路上提供帮助和启发。