人工智能之核心技术 深度学习 第十章 模型部署基础

人工智能之核心技术 深度学习

第十章 模型部署基础

---

@TOC

---

前言：模型部署基础 —— 从训练到生产

“训练是科学，部署是工程。” 模型在实验室跑通只是第一步，高效、稳定、低成本地服务用户才是 AI 价值的真正体现。本章将系统讲解模型保存、优化与部署全流程，并以 文生图 API 服务 为案例，打通端到端部署链路。

---

一、模型保存与加载

1.1 PyTorch：.pth / .pt 文件

保存方式（三种）

import torch
import torch.nn as nn

model = nn.Linear(10, 1)

# 方式1：仅保存参数（推荐）
torch.save(model.state_dict(), 'model.pth')

# 方式2：保存整个模型（不推荐，依赖类定义）
torch.save(model, 'full_model.pth')

# 方式3：保存检查点（含优化器、epoch等）
torch.save({
    'epoch': 10,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'loss': loss,
}, 'checkpoint.pth')

加载方式

# 加载参数（需先实例化模型）
model = nn.Linear(10, 1)
model.load_state_dict(torch.load('model.pth'))
model.eval()  # 切换到推理模式

# 加载完整模型（不推荐）
model = torch.load('full_model.pth')

✅ 最佳实践：

• 永远用 state_dict()：避免 pickle 安全问题和类依赖
• 保存时调用 model.cpu()：确保跨设备兼容
• 命名规范：{task}_{backbone}_{acc}.pth

---

1.2 TensorFlow/Keras：.h5 与 SavedModel（.pb）

保存方式

import tensorflow as tf

model = tf.keras.Sequential([...])

# 方式1：HDF5 格式（仅 Keras 模型）
model.save('model.h5')

# 方式2：SavedModel（推荐，含完整计算图）
model.save('saved_model/')

# 方式3：仅权重
model.save_weights('weights.h5')

加载方式

# 加载 HDF5
model = tf.keras.models.load_model('model.h5')

# 加载 SavedModel（跨语言支持）
model = tf.saved_model.load('saved_model/')

# 加载权重（需先构建相同结构）
model = create_model()
model.load_weights('weights.h5')

📁 SavedModel 目录结构：

saved_model/
├── saved_model.pb      # 计算图
└── variables/
    ├── variables.data-00000-of-00001
    └── variables.index

---

1.3 扩散模型权重加载（Stable Diffusion）

扩散模型通常由 多个子模块 组成，需分别加载：

from diffusers import StableDiffusionPipeline
import torch

# 方法1：从 Hugging Face 自动下载（最简单）
pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16,
    cache_dir="./models"  # 指定本地缓存
)

# 方法2：从本地文件夹加载（已下载）
pipe = StableDiffusionPipeline.from_pretrained(
    "./models--runwayml--stable-diffusion-v1-5/snapshots/...",
    torch_dtype=torch.float16
)

# 方法3：手动加载各组件（高级用法）
from diffusers import AutoencoderKL, UNet2DConditionModel, PNDMScheduler
from transformers import CLIPTextModel, CLIPTokenizer

vae = AutoencoderKL.from_pretrained("path/to/vae")
unet = UNet2DConditionModel.from_pretrained("path/to/unet")
text_encoder = CLIPTextModel.from_pretrained("path/to/text_encoder")
tokenizer = CLIPTokenizer.from_pretrained("path/to/tokenizer")
scheduler = PNDMScheduler.from_pretrained("path/to/scheduler")

pipe = StableDiffusionPipeline(
    vae=vae,
    unet=unet,
    text_encoder=text_encoder,
    tokenizer=tokenizer,
    scheduler=scheduler,
    safety_checker=None,
    feature_extractor=None
)

💡 提示：

• 权重默认存储在 ~/.cache/huggingface/hub/
• 可通过 snapshot_download 预下载：

  from huggingface_hub import snapshot_download
  snapshot_download("runwayml/stable-diffusion-v1-5", local_dir="sd_v1_5")
  

---

二、模型优化：压缩与加速

目标：减小模型体积 + 提升推理速度 + 降低功耗

2.1 模型量化（Quantization）

将 FP32 → INT8，体积减至 1/4，速度提升 2~4 倍。

PyTorch 量化（动态/静态）

# 动态量化（适用于 LSTM/RNN）
quantized_model = torch.quantization.quantize_dynamic(
    model, {nn.Linear}, dtype=torch.qint8
)

# 静态量化（适用于 CNN，需校准数据）
model.eval()
model.qconfig = torch.quantization.get_default_qconfig('fbgemm')
torch.quantization.prepare(model, inplace=True)

# 校准
with torch.no_grad():
    for data, _ in calibration_loader:
        model(data)

# 转换
quantized_model = torch.quantization.convert(model, inplace=True)

TensorFlow 量化（TFLite）

# 转 TFLite 并量化
converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_types = [tf.float16]  # FP16
# 或
converter.representative_dataset = representative_data_gen  # INT8
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]

tflite_quant_model = converter.convert()

⚠️ 注意：

• 量化可能损失精度（<1% 通常可接受）
• INT8 需校准数据集（100~500 张样本）

---

2.2 模型剪枝（Pruning）

移除 不重要的权重（接近 0 的连接）。

TensorFlow 示例（Keras）

import tensorflow_model_optimization as tfmot

# 定义剪枝计划
prune_low_magnitude = tfmot.sparsity.keras.prune_low_magnitude
pruning_params = {
    'pruning_schedule': tfmot.sparsity.keras.PolynomialDecay(
        initial_sparsity=0.30, final_sparsity=0.70,
        begin_step=0, end_step=1000
    )
}

# 应用剪枝
model_for_pruning = prune_low_magnitude(model, **pruning_params)
model_for_pruning.compile(optimizer='adam', loss='sparse_categorical_crossentropy')

# 训练（微调）
model_for_pruning.fit(train_ds, epochs=2)

# 移除剪枝包装，得到紧凑模型
model_for_export = tfmot.sparsity.keras.strip_pruning(model_for_pruning)

📉 效果：模型体积减少 50%~90%，推理速度提升

---

2.3 知识蒸馏（Knowledge Distillation）

用 大模型（Teacher） 指导 小模型（Student） 学习软标签。

graph LR
    A["输入 x"] --> B("Teacher Model")
    A --> C("Student Model")
    B --> D["硬标签 y + 软标签 p_T"]
    C --> E["预测 p_S"]
    D --> F["蒸馏损失：<br>L = α * KL(p_T || p_S) + (1-α) * CE(y, p_S)"]
    E --> F
    F --> G["优化 Student"]

PyTorch 蒸馏损失

import torch.nn.functional as F

def distillation_loss(student_logits, teacher_logits, labels, alpha=0.5, temperature=3):
    # 软标签（带温度）
    soft_targets = F.softmax(teacher_logits / temperature, dim=1)
    soft_prob = F.log_softmax(student_logits / temperature, dim=1)
    
    # KL 散度（蒸馏损失）
    distill_loss = F.kl_div(soft_prob, soft_targets, reduction='batchmean') * (temperature ** 2)
    
    # 交叉熵（任务损失）
    ce_loss = F.cross_entropy(student_logits, labels)
    
    return alpha * distill_loss + (1 - alpha) * ce_loss

✅ 优势：

• Student 模型小、快，精度接近 Teacher
• 适用于移动端部署

---

2.4 扩散模型采样速度优化

Stable Diffusion 默认需 50 步，可通过以下方法加速：

技巧1：使用更快的调度器（Scheduler）

from diffusers import EulerAncestralDiscreteScheduler

pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
# 仅需 20~30 步，质量损失小

技巧2：DDIM / DPM-Solver++

from diffusers import DDIMScheduler, DPMSolverMultistepScheduler

# DDIM：10~20 步即可
pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)

# DPM-Solver++：SOTA 快速采样器
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)

技巧3：TensorRT 加速（见下文）

🚀 效果对比（A100 GPU）：

| 调度器 | 步数 | 时间（秒） | 质量 | |--------|------|----------|------| | PNDM | 50 | 8.2 | ★★★★☆ | | DDIM | 20 | 3.1 | ★★★★ | | DPM++ | 20 | 3.0 | ★★★★★ |

---

三、部署工具实战

3.1 ONNX：跨框架模型转换

ONNX（Open Neural Network Exchange）是 模型中间表示标准，支持 PyTorch → TensorFlow/TensorRT 等。

PyTorch 导出 ONNX

import torch.onnx

# 创建示例输入
dummy_input = torch.randn(1, 3, 224, 224)

# 导出
torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    export_params=True,        # 存储训练参数
    opset_version=11,          # ONNX 算子集版本
    do_constant_folding=True,  # 优化常量
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={
        'input': {0: 'batch_size'},
        'output': {0: 'batch_size'}
    }
)

验证 ONNX 模型

import onnx
onnx_model = onnx.load("model.onnx")
onnx.checker.check_model(onnx_model)  # 验证结构

✅ 优势：

• 一次导出，多平台部署
• 支持 TensorRT、OpenVINO、ONNX Runtime

---

3.2 TensorRT：NVIDIA GPU 极致加速

TensorRT 将模型 编译为 GPU 专属优化引擎，速度提升 3~10 倍。

流程图

graph LR
    A["PyTorch/TensorFlow 模型"] --> B["导出 ONNX"]
    B --> C["TensorRT Builder"]
    C --> D["优化：层融合/精度校准/内核选择"]
    D --> E["TensorRT Engine (.plan)"]
    E --> F["推理 API"]

使用 Python API 构建 Engine

import tensorrt as trt

def build_engine(onnx_file_path, engine_file_path, fp16=True):
    logger = trt.Logger(trt.Logger.WARNING)
    builder = trt.Builder(logger)
    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
    parser = trt.OnnxParser(network, logger)
    
    # 解析 ONNX
    with open(onnx_file_path, 'rb') as model:
        parser.parse(model.read())
    
    config = builder.create_builder_config()
    config.max_workspace_size = 1 << 30  # 1GB
    
    if fp16 and builder.platform_has_fast_fp16:
        config.set_flag(trt.BuilderFlag.FP16)
    
    # 构建 Engine
    engine = builder.build_serialized_network(network, config)
    with open(engine_file_path, "wb") as f:
        f.write(engine)

推理（简化版）

import pycuda.autoinit
import pycuda.driver as cuda
import numpy as np

# 加载 Engine
with open("model.plan", "rb") as f, trt.Runtime(logger) as runtime:
    engine = runtime.deserialize_cuda_engine(f.read())

# 分配内存、执行推理（略）

💡 提示：

• 使用 trtexec 命令行工具更简单：

  trtexec --onnx=model.onnx --saveEngine=model.plan --fp16
  

---

四、实战：搭建文生图 API 服务（FastAPI + Stable Diffusion）

4.1 项目结构

text2image-api/
├── app.py                # FastAPI 主程序
├── model_loader.py       # 模型加载与缓存
├── requirements.txt
└── models/               # 本地模型缓存

4.2 模型加载模块（model_loader.py）

from diffusers import StableDiffusionPipeline
import torch

class ModelManager:
    def __init__(self):
        self.pipe = None
    
    def load_model(self):
        if self.pipe is None:
            self.pipe = StableDiffusionPipeline.from_pretrained(
                "runwayml/stable-diffusion-v1-5",
                torch_dtype=torch.float16,
                cache_dir="./models"
            ).to("cuda")
            self.pipe.enable_attention_slicing()  # 节省内存
        return self.pipe

# 全局单例
model_manager = ModelManager()

4.3 FastAPI 服务（app.py）

from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse
from io import BytesIO
from PIL import Image
from model_loader import model_manager

app = FastAPI(title="Text-to-Image API")

@app.post("/generate")
async def generate_image(prompt: str, num_inference_steps: int = 20):
    try:
        pipe = model_manager.load_model()
        
        # 生成图像
        image = pipe(
            prompt,
            num_inference_steps=num_inference_steps,
            guidance_scale=7.5
        ).images[0]
        
        # 转为字节流
        img_io = BytesIO()
        image.save(img_io, 'PNG')
        img_io.seek(0)
        
        return StreamingResponse(img_io, media_type="image/png")
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)

4.4 启动与测试

# 安装依赖
pip install fastapi uvicorn diffusers transformers accelerate

# 启动服务
python app.py

# 测试（命令行）
curl -X POST "http://localhost:8000/generate" \
     -H "Content-Type: application/json" \
     -d '{"prompt":"a red apple on a white table"}' \
     --output apple.png

⚙️ 生产优化建议：

• 使用 异步队列（Celery + Redis）处理高并发
• 添加 速率限制（SlowAPI）
• 部署到 Docker 容器
• 使用 TensorRT 加速（替换 Diffusers 推理后端）

---

五、总结与最佳实践

阶段	关键技术	工具
保存	state_dict / SavedModel	PyTorch, TF
优化	量化/剪枝/蒸馏	TFLite, TorchQuant
转换	ONNX 中间表示	torch.onnx
加速	TensorRT / OpenVINO	NVIDIA, Intel
部署	REST API / gRPC	FastAPI, Flask

🔚 部署 Checklist：

1. ✅ 模型是否已转为推理模式？（model.eval() / training=False）
2. ✅ 是否进行了量化/剪枝？
3. ✅ 是否使用 ONNX/TensorRT 加速？
4. ✅ API 是否有错误处理和限流？
5. ✅ 是否监控延迟与 GPU 利用率？

“部署不是终点，而是价值的起点。” 掌握本章技能，你已具备将 AI 模型转化为真实产品的能力！

---

附录：常用命令

# 导出 ONNX
python -c "import torch; model=torch.load('model.pth'); torch.onnx.export(...)"

# TensorRT 构建 Engine
trtexec --onnx=model.onnx --saveEngine=model.plan --fp16

# 启动 FastAPI
uvicorn app:app --reload --port 8000

# Dockerfile 示例
FROM nvidia/cuda:12.1-base
RUN pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
RUN pip install diffusers transformers accelerate fastapi uvicorn
COPY . /app
WORKDIR /app
CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "8000"]

资料关注

公众号：咚咚王 gitee：gitee.com/wy185850518…

《Python编程：从入门到实践》 《利用Python进行数据分析》 《算法导论中文第三版》 《概率论与数理统计（第四版） (盛骤) 》 《程序员的数学》 《线性代数应该这样学第3版》 《微积分和数学分析引论》 《（西瓜书）周志华-机器学习》 《TensorFlow机器学习实战指南》 《Sklearn与TensorFlow机器学习实用指南》 《模式识别（第四版）》 《深度学习 deep learning》伊恩·古德费洛著 花书 《Python深度学习第二版(中文版)【纯文本】 (登封大数据 (Francois Choliet)) (Z-Library)》 《深入浅出神经网络与深度学习+(迈克尔·尼尔森（Michael+Nielsen）》 《自然语言处理综论 第2版》 《Natural-Language-Processing-with-PyTorch》 《计算机视觉-算法与应用(中文版)》 《Learning OpenCV 4》 《AIGC：智能创作时代》杜雨+&+张孜铭 《AIGC原理与实践：零基础学大语言模型、扩散模型和多模态模型》 《从零构建大语言模型（中文版）》 《实战AI大模型》 《AI 3.0》