4. 模型准备与优化：模型格式转换、量化与剪枝

在推理服务部署链路中，模型从训练框架产出到生产环境可用，通常不是“拷贝权重文件”那么简单。训练框架（如 PyTorch、TensorFlow）优先考虑迭代灵活性，而推理引擎（如 TensorRT、ONNX Runtime、OpenVINO）则追求计算图固定、算子融合和低精度计算。两者之间存在明显的“语义鸿沟”。本章聚焦部署前的三道关键工序：格式转换、量化与剪枝，帮助你将一个“可训练的模型”转变为“可高效运行的推理模型”。

4.1 核心原理：为什么不能直接部署训练模型？

训练框架保存的模型（如 .pt、.h5）本质上是动态计算图的序列化结果。动态图在每次前向传播时都会重新构建算子执行路径，这为研究员调试提供了便利，但给推理引擎带来了三个致命问题：

  • 算子粒度太细：PyTorch 中一个简单的 LayerNorm 可能被拆解为 10+ 个小算子，推理引擎无法有效融合，导致 GPU 内核启动开销过大。
  • 隐式控制流：动态 if 或 for 循环依赖数据值，导致推理引擎无法静态规划显存和线程块。
  • 精度冗余：训练时使用 FP32 以保证梯度稳定性，但推理时多数场景并不需要如此高的精度。

因此，部署前的优化本质上是回答三个问题：用什么格式表达计算图？用什么数值精度执行计算？用什么结构减少计算量？ 下面分别展开。

核心结论：格式转换解决“引擎看不懂”的问题；量化解决“算得太慢”的问题；剪枝解决“算得太多”的问题。三者顺序通常为：先导出通用格式（ONNX），再针对特定引擎优化（TensorRT），最后根据硬件能力决定是否量化或剪枝。

4.2 模型格式转换：从 PyTorch 到 ONNX 再到 TensorRT

ONNX（Open Neural Network Exchange）是目前最通用的中间表示。它定义了一套静态计算图规范，几乎所有推理引擎都能消费。转换的核心操作是追踪（tracing）或脚本化（scripting）。

以 PyTorch 为例，最常用的导出代码如下：

import torch
import torch.onnx

model = MyModel().eval()
dummy_input = torch.randn(1, 3, 224, 224)

torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    opset_version=17,
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}}
)
print("ONNX 导出完成")

导出后，建议用 onnxruntime 或 onnx.checker 验证图完整性。若遇到算子不支持，通常需要算子融合或自定义算子。例如，将 torch.where 替换为 ONNX Where 节点，或将多个小算子合并为单个 Gemm。

当目标硬件为 NVIDIA GPU 时，ONNX 通常只是“中间站”，最终需要转换为 TensorRT 引擎。TensorRT 的转换工具 trtexec 提供了命令行方式：

trtexec --onnx=model.onnx \
        --saveEngine=model.engine \
        --fp16 \
        --minShapes=input:1x3x224x224 \
        --optShapes=input:8x3x224x224 \
        --maxShapes=input:16x3x224x224

上述命令启用了 FP16 精度，并设置了动态 batch 范围。TensorRT 会执行层融合（如 Conv+BN+ReLU 融合为 CBR）、内核自动调优（选择最优 CUDA kernel）以及显存池复用。转换后的 .engine 文件是二进制序列化格式，只能在相同 GPU 架构和 TensorRT 版本上反序列化。

格式 适用场景 优点 缺点
PyTorch (.pt) 训练、调试 灵活、动态图 推理慢、依赖 Python
ONNX 跨平台中间格式 通用、静态图 算子覆盖有限
TensorRT (.engine) NVIDIA GPU 生产 极致性能、低延迟 硬件绑定、编译时间长
OpenVINO IR Intel CPU/GPU CPU 优化好 NVIDIA 支持弱

4.3 量化：用更少的比特表达权重与激活

量化是将 FP32 数值映射到低比特（如 INT8、INT4）的过程。其核心收益在于：内存带宽减半或减至 1/4，同时利用硬件 INT8 算力（通常为 FP32 的 2~4 倍）。但量化会引入信息损失，因此需要校准策略。

量化分为两类：

  • 训练后量化（PTQ）：无需重新训练，只需提供少量校准数据（通常 500~2000 张），统计每层激活值的动态范围，计算缩放因子 scale 和零点 zero_point。适合快速部署。
  • 量化感知训练（QAT）：在训练过程中模拟量化误差，让模型权重适应低比特表示。通常比 PTQ 精度高 1~3%，但需要完整训练流程。

以 TensorRT 的 PTQ 为例，需要准备校准数据并生成校准缓存：

import tensorrt as trt

def calibrator():
    # 自定义 Calibrator 类，读取校准图片
    return MyCalibrator(calib_images, batch_size=8)

builder = trt.Builder(TRT_LOGGER)
network = builder.create_network()
parser = trt.OnnxParser(network, TRT_LOGGER)
with open("model.onnx", "rb") as f:
    parser.parse(f.read())

config = builder.create_builder_config()
config.set_flag(trt.BuilderFlag.INT8)
config.int8_calibrator = calibrator()
engine = builder.build_engine(network, config)

量化后务必评估精度。通常使用验证集计算 mAP 或准确率，若下降超过 1%，建议回退到 FP16 或改用 QAT。此外，敏感层分析是常用技巧：逐层量化并观察精度变化，找出对量化敏感的层（如检测头的回归分支），对这些层保持 FP16。

实操建议：对于视觉模型，优先尝试 PTQ + INT8。若精度不达标，再考虑 QAT。对于 NLP 模型（如 BERT），注意 Embedding 层和 LayerNorm 层通常对量化敏感，建议保留 FP16。

4.4 剪枝：结构化与非结构化

剪枝是移除模型中冗余权重或神经元的过程。它直接减少计算量（FLOPs）和模型体积。但剪枝需要与硬件特性配合，否则可能“剪了白剪”。

剪枝分为两类：

  • 非结构化剪枝：将权重矩阵中绝对值较小的元素置零。稀疏度高，但产生不规则稀疏矩阵，GPU 上难以加速（除非使用专用稀疏内核，如 NVIDIA 2:4 结构化稀疏）。
  • 结构化剪枝：按通道（Channel）或滤波器（Filter）整体移除。保持矩阵稠密规则，可直接减少计算量，对硬件友好。

以 PyTorch 的 torch.nn.utils.prune 为例，对卷积层进行通道剪枝的伪代码如下：

import torch.nn.utils.prune as prune

# 对 Conv2d 的权重按 L1 范数剪枝 20% 通道
prune.l1_unstructured(conv_layer, name="weight", amount=0.2)
# 剪枝后需要重新打包为结构化稀疏
prune.remove(conv_layer, "weight")

更工程化的做法是使用 Intel Neural Compressor 或 NVIDIA TensorRT Model Optimizer，它们支持自动化剪枝与量化联合优化。剪枝后必须进行微调（fine-tune），通常需要 1~2 个 epoch 恢复精度。

4.5 案例：YOLOv8 检测模型部署优化全流程

假设我们有一个训练好的 YOLOv8s 模型（COCO 80 类），目标部署到 NVIDIA A10 GPU，要求延迟 < 5ms。

  1. 导出 ONNX：使用 yolo export model=yolov8s.pt format=onnx dynamic=True，得到动态 batch 的 ONNX 文件。
  2. TensorRT 转换：执行 trtexec --onnx=yolov8s.onnx --saveEngine=yolov8s.engine --fp16，测得 FP16 延迟约 3.2ms，满足要求。
  3. 尝试 INT8 量化：准备 1000 张 COCO 验证集子集作为校准数据，执行 INT8 PTQ。量化后延迟降至 2.1ms，但 mAP@0.5 从 0.52 降至 0.49，下降 5.7%，不可接受。
  4. 敏感层分析：发现检测头（Detect Head）的最后一层卷积对量化敏感。将该层保留 FP16，其余层 INT8。重新构建引擎，mAP 恢复至 0.51，延迟 2.3ms。
  5. 结构化剪枝：对 Backbone 的 C2f 模块按通道 L1 范数剪枝 10%，微调 10 个 epoch，mAP 回升至 0.515，延迟降至 2.0ms。

最终部署方案：TensorRT INT8 + 部分 FP16 混合精度 + Backbone 10% 通道剪枝，在满足精度损失 <1% 的前提下，延迟从 FP32 的 6.8ms 降至 2.0ms，加速 3.4 倍。

4.6 常见误区与复盘

误区一：量化后不评估就直接上线。量化误差在极端输入下可能放大，必须用与线上分布一致的验证集评估，并设置精度回退开关。

误区二：剪枝只看 FLOPs 减少。FLOPs 减少 50% 不代表延迟减半。内存访问瓶颈、kernel launch 开销可能占主导。务必在目标硬件上实测延迟。

误区三：ONNX 导出后不做算子兼容检查。某些 PyTorch 算子（如 torch.fft）在 ONNX 中无对应实现，导出时静默失败或产生错误节点。建议导出后使用 onnxruntime 跑一遍推理验证输出一致性。

误区四：TensorRT 引擎跨设备复用。TensorRT 引擎与 GPU 架构（如 Ampere、Ada）、CUDA 版本、TensorRT 版本强绑定。换机器必须重新构建引擎，否则报错或性能骤降。

复盘一个真实事故：某团队将 TensorRT FP16 引擎从 A100 迁移到 V100，未重新构建，结果推理结果全部错误。原因是 A100 支持 TF32 而 V100 不支持，导致某些算子回退到 FP32，数值行为改变。正确做法是：在每台目标机器上执行一次构建流程，并将引擎构建纳入 CI/CD 流水线。

另一个常见问题是校准数据与线上数据分布不一致。某 OCR 模型使用通用文本图片校准，上线后遇到大量票据图片，激活值范围远超校准集，导致 INT8 精度崩溃。解决方案是使用线上真实采样数据（至少 1000 张）重新校准，并定期更新校准缓存。

最后，优化并非一劳永逸。模型结构、硬件驱动、推理引擎版本都在演进。建议建立基准测试脚本，每次变更后自动跑精度与延迟回归，确保优化不引入隐性退化。

模型部署优化流水线 PyTorch 模型 动态图 · FP32 .pt / .pth 导出 ONNX 中间格式 静态图 · 通用 model.onnx 转换 TensorRT 引擎 算子融合 · 显存优化 model.engine 生产推理服务 GPU / CPU 低延迟 · 高吞吐 优化手段（可叠加） 量化 (Quantization) FP32 → INT8 / FP16 PTQ：无需训练，需校准集 QAT：训练中模拟量化误差 收益：带宽↓ 2~4x 算力↑ 剪枝 (Pruning) 移除冗余权重/通道 非结构化：稀疏矩阵 结构化：按通道/滤波器 收益：FLOPs↓ 体积↓ 混合精度策略 敏感层保留 FP16 检测头/Embedding 层 其余层 INT8 收益：精度与速度平衡 评估闭环：精度验证 → 延迟实测 → 回退/调参 mAP / 准确率 < 阈值？→ 调整量化层或剪枝率 → 重新构建 推荐顺序：PyTorch → ONNX → TensorRT → (量化 + 剪枝) → 精度/延迟评估 注意：TensorRT 引擎与 GPU 架构强绑定，换机需重新构建

上图总结了本章的核心流程。从 PyTorch 动态图出发，经过 ONNX 标准化，再到 TensorRT 引擎优化，最后叠加量化与剪枝手段，并始终围绕“精度-延迟”闭环进行评估。每一步都有对应的工具链和验证方法，切忌跳过评估直接上线。

← 上一章 📖 返回目录 下一章 →