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

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

第十章 模型部署基础

文章目录

• 人工智能之核心技术 深度学习
• 前言：模型部署基础 —— 从训练到生产
• • 一、模型保存与加载
  • • 1.1 PyTorch：`.pth` / `.pt` 文件
    • • 保存方式（三种）
      • 加载方式
    • 1.2 TensorFlow/Keras：`.h5` 与 SavedModel（`.pb`）
    • • 保存方式
      • 加载方式
    • 1.3 扩散模型权重加载（Stable Diffusion）
  • 二、模型优化：压缩与加速
  • • 2.1 模型量化（Quantization）
    • • PyTorch 量化（动态/静态）
      • TensorFlow 量化（TFLite）
    • 2.2 模型剪枝（Pruning）
    • • TensorFlow 示例（Keras）
    • 2.3 知识蒸馏（Knowledge Distillation）
    • • PyTorch 蒸馏损失
    • 2.4 扩散模型采样速度优化
    • • 技巧1：使用更快的调度器（Scheduler）
      • 技巧2：DDIM / DPM-Solver++
      • 技巧3：TensorRT 加速（见下文）
  • 三、部署工具实战
  • • 3.1 ONNX：跨框架模型转换
    • • PyTorch 导出 ONNX
      • 验证 ONNX 模型
    • 3.2 TensorRT：NVIDIA GPU 极致加速
    • • 流程图
      • 使用 Python API 构建 Engine
      • 推理（简化版）
  • 四、实战：搭建文生图 API 服务（FastAPI + Stable Diffusion）
  • • 4.1 项目结构
    • 4.2 模型加载模块（`model_loader.py`）
    • 4.3 FastAPI 服务（`app.py`）
    • 4.4 启动与测试
  • 五、总结与最佳实践
• 资料关注

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

“训练是科学，部署是工程。” 模型在实验室跑通只是第一步，高效、稳定、低成本地服务用户才是 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） 学习软标签。

#mermaid-svg-DftjVktojDdi0hi4{font-family:\”trebuchet ms\”,verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-DftjVktojDdi0hi4 .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-DftjVktojDdi0hi4 .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-DftjVktojDdi0hi4 .error-icon{fill:#552222;}#mermaid-svg-DftjVktojDdi0hi4 .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-DftjVktojDdi0hi4 .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-DftjVktojDdi0hi4 .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-DftjVktojDdi0hi4 .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-DftjVktojDdi0hi4 .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-DftjVktojDdi0hi4 .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-DftjVktojDdi0hi4 .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-DftjVktojDdi0hi4 .marker{fill:#333333;stroke:#333333;}#mermaid-svg-DftjVktojDdi0hi4 .marker.cross{stroke:#333333;}#mermaid-svg-DftjVktojDdi0hi4 svg{font-family:\”trebuchet ms\”,verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-DftjVktojDdi0hi4 p{margin:0;}#mermaid-svg-DftjVktojDdi0hi4 .label{font-family:\”trebuchet ms\”,verdana,arial,sans-serif;color:#333;}#mermaid-svg-DftjVktojDdi0hi4 .cluster-label text{fill:#333;}#mermaid-svg-DftjVktojDdi0hi4 .cluster-label span{color:#333;}#mermaid-svg-DftjVktojDdi0hi4 .cluster-label span p{background-color:transparent;}#mermaid-svg-DftjVktojDdi0hi4 .label text,#mermaid-svg-DftjVktojDdi0hi4 span{fill:#333;color:#333;}#mermaid-svg-DftjVktojDdi0hi4 .node rect,#mermaid-svg-DftjVktojDdi0hi4 .node circle,#mermaid-svg-DftjVktojDdi0hi4 .node ellipse,#mermaid-svg-DftjVktojDdi0hi4 .node polygon,#mermaid-svg-DftjVktojDdi0hi4 .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-DftjVktojDdi0hi4 .rough-node .label text,#mermaid-svg-DftjVktojDdi0hi4 .node .label text,#mermaid-svg-DftjVktojDdi0hi4 .image-shape .label,#mermaid-svg-DftjVktojDdi0hi4 .icon-shape .label{text-anchor:middle;}#mermaid-svg-DftjVktojDdi0hi4 .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-DftjVktojDdi0hi4 .rough-node .label,#mermaid-svg-DftjVktojDdi0hi4 .node .label,#mermaid-svg-DftjVktojDdi0hi4 .image-shape .label,#mermaid-svg-DftjVktojDdi0hi4 .icon-shape .label{text-align:center;}#mermaid-svg-DftjVktojDdi0hi4 .node.clickable{cursor:pointer;}#mermaid-svg-DftjVktojDdi0hi4 .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-DftjVktojDdi0hi4 .arrowheadPath{fill:#333333;}#mermaid-svg-DftjVktojDdi0hi4 .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-DftjVktojDdi0hi4 .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-DftjVktojDdi0hi4 .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-DftjVktojDdi0hi4 .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-DftjVktojDdi0hi4 .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-DftjVktojDdi0hi4 .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-DftjVktojDdi0hi4 .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-DftjVktojDdi0hi4 .cluster text{fill:#333;}#mermaid-svg-DftjVktojDdi0hi4 .cluster span{color:#333;}#mermaid-svg-DftjVktojDdi0hi4 div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:\”trebuchet ms\”,verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-DftjVktojDdi0hi4 .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-DftjVktojDdi0hi4 rect.text{fill:none;stroke-width:0;}#mermaid-svg-DftjVktojDdi0hi4 .icon-shape,#mermaid-svg-DftjVktojDdi0hi4 .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-DftjVktojDdi0hi4 .icon-shape p,#mermaid-svg-DftjVktojDdi0hi4 .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-DftjVktojDdi0hi4 .icon-shape rect,#mermaid-svg-DftjVktojDdi0hi4 .image-shape rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-DftjVktojDdi0hi4 .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-DftjVktojDdi0hi4 .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-DftjVktojDdi0hi4 :root{–mermaid-font-family:\”trebuchet ms\”,verdana,arial,sans-serif;}

输入 x

Teacher Model

Student Model

硬标签 y + 软标签 p_T

预测 p_S

蒸馏损失：L = α * KL(p_T || p_S) + (1-α) * CE(y, p_S)

优化 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 倍。

流程图

#mermaid-svg-m8BRntrZArOnfU2c{font-family:\”trebuchet ms\”,verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-m8BRntrZArOnfU2c .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-m8BRntrZArOnfU2c .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-m8BRntrZArOnfU2c .error-icon{fill:#552222;}#mermaid-svg-m8BRntrZArOnfU2c .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-m8BRntrZArOnfU2c .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-m8BRntrZArOnfU2c .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-m8BRntrZArOnfU2c .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-m8BRntrZArOnfU2c .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-m8BRntrZArOnfU2c .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-m8BRntrZArOnfU2c .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-m8BRntrZArOnfU2c .marker{fill:#333333;stroke:#333333;}#mermaid-svg-m8BRntrZArOnfU2c .marker.cross{stroke:#333333;}#mermaid-svg-m8BRntrZArOnfU2c svg{font-family:\”trebuchet ms\”,verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-m8BRntrZArOnfU2c p{margin:0;}#mermaid-svg-m8BRntrZArOnfU2c .label{font-family:\”trebuchet ms\”,verdana,arial,sans-serif;color:#333;}#mermaid-svg-m8BRntrZArOnfU2c .cluster-label text{fill:#333;}#mermaid-svg-m8BRntrZArOnfU2c .cluster-label span{color:#333;}#mermaid-svg-m8BRntrZArOnfU2c .cluster-label span p{background-color:transparent;}#mermaid-svg-m8BRntrZArOnfU2c .label text,#mermaid-svg-m8BRntrZArOnfU2c span{fill:#333;color:#333;}#mermaid-svg-m8BRntrZArOnfU2c .node rect,#mermaid-svg-m8BRntrZArOnfU2c .node circle,#mermaid-svg-m8BRntrZArOnfU2c .node ellipse,#mermaid-svg-m8BRntrZArOnfU2c .node polygon,#mermaid-svg-m8BRntrZArOnfU2c .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-m8BRntrZArOnfU2c .rough-node .label text,#mermaid-svg-m8BRntrZArOnfU2c .node .label text,#mermaid-svg-m8BRntrZArOnfU2c .image-shape .label,#mermaid-svg-m8BRntrZArOnfU2c .icon-shape .label{text-anchor:middle;}#mermaid-svg-m8BRntrZArOnfU2c .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-m8BRntrZArOnfU2c .rough-node .label,#mermaid-svg-m8BRntrZArOnfU2c .node .label,#mermaid-svg-m8BRntrZArOnfU2c .image-shape .label,#mermaid-svg-m8BRntrZArOnfU2c .icon-shape .label{text-align:center;}#mermaid-svg-m8BRntrZArOnfU2c .node.clickable{cursor:pointer;}#mermaid-svg-m8BRntrZArOnfU2c .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-m8BRntrZArOnfU2c .arrowheadPath{fill:#333333;}#mermaid-svg-m8BRntrZArOnfU2c .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-m8BRntrZArOnfU2c .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-m8BRntrZArOnfU2c .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-m8BRntrZArOnfU2c .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-m8BRntrZArOnfU2c .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-m8BRntrZArOnfU2c .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-m8BRntrZArOnfU2c .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-m8BRntrZArOnfU2c .cluster text{fill:#333;}#mermaid-svg-m8BRntrZArOnfU2c .cluster span{color:#333;}#mermaid-svg-m8BRntrZArOnfU2c div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:\”trebuchet ms\”,verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-m8BRntrZArOnfU2c .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-m8BRntrZArOnfU2c rect.text{fill:none;stroke-width:0;}#mermaid-svg-m8BRntrZArOnfU2c .icon-shape,#mermaid-svg-m8BRntrZArOnfU2c .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-m8BRntrZArOnfU2c .icon-shape p,#mermaid-svg-m8BRntrZArOnfU2c .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-m8BRntrZArOnfU2c .icon-shape rect,#mermaid-svg-m8BRntrZArOnfU2c .image-shape rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-m8BRntrZArOnfU2c .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-m8BRntrZArOnfU2c .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-m8BRntrZArOnfU2c :root{–mermaid-font-family:\”trebuchet ms\”,verdana,arial,sans-serif;}

PyTorch/TensorFlow 模型

导出 ONNX

TensorRT Builder

优化：层融合/精度校准/内核选择

TensorRT Engine (.plan)

推理 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 推理后端）

五、总结与最佳实践

🔚 部署 Checklist：

“部署不是终点，而是价值的起点。” 掌握本章技能，你已具备将 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：https://gitee.com/wy18585051844/ai_learning

《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》