📚 课程目录第4章模型打包与容器化

4. 模型打包与容器化：使用 Docker 打包 PyTorch/TensorFlow 模型

模型打包与容器化是连接训练环境与生产推理环境的桥梁。本章将深入讲解如何将训练好的 PyTorch/TensorFlow 模型通过 Docker 容器化，构建标准化的模型仓库结构，并实现多版本管理，确保模型在在线推理链路中可重复、可隔离、可扩展。

4.1 模型打包的核心目标与挑战

在进入具体操作前，需要明确模型打包需要解决的核心问题：

环境一致性：消除“在我机器上能跑”的魔咒，锁定 CUDA、cuDNN、Python 依赖版本。
依赖隔离：不同模型可能依赖不同版本的框架或库，容器化提供进程级隔离。
启动效率：打包后的镜像应支持秒级加载模型权重，避免冷启动过慢。
资源可观测：容器化后便于通过 K8s 等编排工具进行 GPU/内存资源的分配与监控。

实践建议：生产环境中，模型打包不应仅包含模型权重文件（.pt/.pth/.h5），还应包含预处理逻辑（tokenizer、图像归一化参数）和后处理逻辑（阈值过滤、结果排序），确保推理链路的完整性。

4.2 模型仓库结构（Model Repository）

一个标准化的模型仓库结构是支撑多版本管理和高效部署的基础。推荐采用以下目录规范：

层级	目录/文件	说明
根目录	`model_repository/`	模型仓库根目录，可包含多个模型
模型名称	`sentiment_model/`	模型唯一标识，与推理服务配置中的模型名对应
版本目录	`1/`, `2/`, `3/`	整数版本号，Triton Inference Server 等工具按此加载
模型文件	`model.pt` / `model.pth`	序列化的模型权重文件（PyTorch 常用）
配置文件	`config.pbtxt`	模型元信息：输入输出张量名称、数据类型、形状、批处理策略
预处理	`tokenizer.json` / `preprocess.py`	分词器或预处理脚本，可选但强烈推荐
标签映射	`labels.txt`	类别索引到实际标签的映射文件

  关键点：版本目录使用纯数字命名（如 1, 2），而非语义化版本（如 v1.0.0）。推理服务框架（如 Triton、TorchServe）通过轮询版本目录的变化来实现模型的平滑升级与回滚。

以下是一个典型的模型仓库结构示意图，展示了多模型、多版本的层级关系：

4.3 使用 Docker 打包 PyTorch 模型

以 PyTorch 模型为例，展示从模型导出到 Docker 镜像构建的完整流程。核心思路是将模型权重、推理脚本和依赖环境打包成一个可独立运行的容器。

4.3.1 模型导出与序列化

推荐使用 torch.jit.script 或 torch.jit.trace 将模型导出为 TorchScript 格式，消除对原始 Python 类定义的依赖：

import torch
import torchvision.models as models

# 加载预训练模型
model = models.resnet18(pretrained=True)
model.eval()

# 使用 TorchScript trace 导出
dummy_input = torch.randn(1, 3, 224, 224)
traced_model = torch.jit.trace(model, dummy_input)

# 保存为 model.pt
traced_model.save("model_repository/image_classifier/1/model.pt")
print("模型已导出为 TorchScript 格式")

4.3.2 编写推理服务脚本

创建一个 inference.py 文件，包含模型加载、预处理、推理和后处理逻辑：

import torch
import torchvision.transforms as transforms
from PIL import Image
import json

class ModelInference:
    def __init__(self, model_path, device="cuda"):
        self.device = torch.device(device if torch.cuda.is_available() else "cpu")
        self.model = torch.jit.load(model_path, map_location=self.device)
        self.model.eval()
        
        # 预处理流水线
        self.transform = transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                 std=[0.229, 0.224, 0.225])
        ])
        
        # 加载标签映射
        with open("labels.txt", "r") as f:
            self.labels = [line.strip() for line in f.readlines()]
    
    def predict(self, image_bytes):
        # 图像解码与预处理
        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
        input_tensor = self.transform(image).unsqueeze(0).to(self.device)
        
        # 推理
        with torch.no_grad():
            output = self.model(input_tensor)
        
        # 后处理：获取 Top-5 结果
        probabilities = torch.nn.functional.softmax(output[0], dim=0)
        top5_prob, top5_indices = torch.topk(probabilities, 5)
        
        results = []
        for i in range(5):
            idx = top5_indices[i].item()
            results.append({
                "label": self.labels[idx],
                "confidence": round(top5_prob[i].item(), 4)
            })
        return results

4.3.3 编写 Dockerfile

采用多阶段构建策略，减小最终镜像体积：

# ===== 第一阶段：基础环境 =====
FROM pytorch/pytorch:2.1.0-cuda12.1-cudnn8-runtime AS base

# 设置工作目录
WORKDIR /app

# 安装系统依赖
RUN apt-get update && apt-get install -y --no-install-recommends \
    libgl1-mesa-glx \
    libglib2.0-0 \
    && rm -rf /var/lib/apt/lists/*

# 安装 Python 依赖
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# ===== 第二阶段：生产镜像 =====
FROM base AS production

# 复制模型仓库
COPY model_repository/ ./model_repository/

# 复制推理服务脚本
COPY inference.py .
COPY server.py .

# 暴露推理端口
EXPOSE 8080

# 启动命令
CMD ["python", "server.py", "--model-path", "./model_repository/image_classifier/1/model.pt"]

注意：基础镜像选择需与训练环境保持一致。如果训练使用 PyTorch 2.1.0 + CUDA 12.1，则推理镜像也应使用相同版本，避免算子兼容性问题。切勿随意升级或降级 CUDA 版本。

4.3.4 构建与运行

# 构建镜像
docker build -t image-classifier:v1.0.0 -f Dockerfile .

# 运行容器（GPU 模式）
docker run --gpus all -p 8080:8080 \
  -v $(pwd)/model_repository:/app/model_repository \
  image-classifier:v1.0.0

# 验证推理
curl -X POST http://localhost:8080/predict \
  -F "image=@test.jpg"

4.4 使用 Docker 打包 TensorFlow 模型

TensorFlow 模型的打包流程与 PyTorch 类似，但导出格式和推理框架有所不同。推荐使用 SavedModel 格式：

import tensorflow as tf

# 加载训练好的模型
model = tf.keras.models.load_model("saved_model/my_model")

# 导出为 SavedModel 格式
model.save("model_repository/tf_classifier/1/", save_format="tf")
print("模型已导出为 SavedModel 格式")

TensorFlow Serving 的 Dockerfile 示例：

FROM tensorflow/serving:2.13.0-gpu

# 复制模型仓库
COPY model_repository/ /models/

# 设置模型配置
ENV MODEL_NAME=tf_classifier
ENV MODEL_BASE_PATH=/models/tf_classifier

# 暴露 gRPC 和 REST 端口
EXPOSE 8500 8501

# 启动 TensorFlow Serving
CMD ["--model_config_file=/models/models.config", "--monitoring_config_file=/models/monitoring.config"]

4.5 多版本管理策略

多版本管理是模型迭代的核心能力。以下表格对比了三种常见的版本管理策略：

策略	实现方式	优势	劣势	适用场景
目录版本	模型仓库中创建 `1/`, `2/`, `3/` 子目录	结构简单，Triton/TF Serving 原生支持	版本回滚需手动调整目录符号链接	中小规模团队，版本迭代频率低
镜像标签	Docker 镜像使用语义化标签 `v1.0.0`, `v2.0.0`	完全隔离，环境独立，回滚只需切换镜像	镜像体积大，拉取时间长，存储成本高	需要严格环境隔离的生产场景
配置路由	通过推理服务配置（如 Istio + K8s）实现流量灰度	支持 A/B 测试、金丝雀发布、蓝绿部署	运维复杂度高，需要服务网格支持	大规模在线推理，需要精细化流量管理

以下示意图展示了基于镜像标签的多版本管理流程：

4.6 最佳实践与避坑指南

  最佳实践清单：
  最小化镜像体积：使用多阶段构建，仅保留运行时依赖。PyTorch 镜像可从 3GB 压缩至 800MB。
固定基础镜像版本：使用 pytorch:2.1.0-cuda12.1 而非 pytorch:latest，避免意外升级导致兼容性问题。
模型预热：在容器启动脚本中加入模型预热步骤，加载权重并执行一次虚拟推理，避免首次请求超时。
健康检查：在 Dockerfile 中添加 HEALTHCHECK 指令，定期检测推理服务是否正常响应。
日志结构化：推理日志输出为 JSON 格式，便于 ELK/Loki 等日志系统采集和分析。

常见陷阱：

CUDA 版本不匹配：训练时使用 CUDA 11.8，推理镜像使用 CUDA 12.1，导致算子编译失败。务必保持 CUDA 运行时版本一致。
文件权限问题：容器内运行用户非 root 时，模型文件权限需设置为 644 或 755，否则加载失败。
忽略 .dockerignore：未添加 .dockerignore 文件会导致将本地 __pycache__、.git 等无用文件打包进镜像，增加构建时间和镜像体积。
单进程模型加载：多线程推理时，模型加载需加锁，避免多个 worker 同时加载导致 OOM。

4.7 本章小结

模型打包与容器化是 AML 在线推理链路中承上启下的关键环节。通过本章学习，应掌握：

标准模型仓库结构的搭建方法，支持多模型、多版本共存。
PyTorch 和 TensorFlow 模型的 Docker 打包流程，包括模型导出、Dockerfile 编写、多阶段构建。
三种多版本管理策略的优劣对比，以及基于镜像标签的版本管理流程。
生产环境中的最佳实践和常见陷阱，确保容器化后的模型推理服务稳定、高效。