📚 课程目录第4章 ONNX模型转换与优化

4、ONNX模型转换与优化：从PyTorch/TensorFlow到TensorRT Engine

ONNX（Open Neural Network Exchange）是连接训练框架与推理引擎的开放中间表示。在TensorRT服务化部署链路中，模型能否高效、无损地从PyTorch/TensorFlow迁移到TensorRT，取决于ONNX导出质量、图优化程度以及Engine生成策略。本章围绕“导出—校验—优化—生成”四个环节，系统讲解ONNX转换的完整方法论，并结合trtexec与Python API给出可落地的操作路径。

4.1 ONNX导出的核心原理与操作流程

ONNX导出的本质是将训练框架的计算图映射为ONNX算子序列，并保留张量的数据类型、形状信息以及控制流语义。PyTorch通过torch.onnx.export采用追踪（Tracing）机制，即实际运行一次模型并记录执行路径；TensorFlow则通过tf2onnx.convert将GraphDef或SavedModel转换为ONNX。两种方式各有局限：追踪无法捕获动态控制流（如if依赖张量值），而静态图转换则可能丢失部分训练期信息。

关键导出参数（PyTorch示例）：

import torch
  import torchvision.models as models
  
  model = models.resnet18(pretrained=True).eval()
  dummy_input = torch.randn(1, 3, 224, 224)
  
  torch.onnx.export(
      model, dummy_input, "resnet18.onnx",
      opset_version=17,                # 推荐≥11，17支持更多动态算子
      input_names=["input"],           # 固定名称便于TensorRT绑定
      output_names=["output"],
      dynamic_axes={
          "input": {0: "batch_size"},  # 仅batch维度动态
          "output": {0: "batch_size"}
      },
      do_constant_folding=True,        # 常量折叠优化
      export_params=True               # 导出权重
  )

导出后必须进行ONNX模型验证。使用onnx.checker.check_model检查结构合法性，再用onnxruntime进行数值一致性验证。建议对比PyTorch输出与ONNX Runtime输出的余弦相似度，阈值设为0.999以上。若出现算子不支持，优先考虑升级opset版本或修改模型结构（如将torch.where替换为where算子）。

检查项	工具/方法	通过标准	失败处理
结构合法性	onnx.checker.check_model	无异常	检查算子版本、输入输出定义
数值一致性	onnxruntime + numpy对比	余弦相似度≥0.999	检查预处理差异、算子精度
动态轴正确性	修改batch=4/8重新推理	输出shape正确	重新导出并标记dynamic_axes
算子兼容性	onnxruntime 或 polygraphy	无unsupported算子	替换算子或升级opset

4.2 动态形状设置与精度校准策略

TensorRT要求Engine在构建时确定张量形状范围。动态形状通过minShapes、optShapes、maxShapes三个维度定义，分别对应最小、最优、最大输入尺寸。最优形状直接影响Engine的kernel选择，应设置为实际部署中最常见的batch size和分辨率。例如，视频分析服务通常batch=1，而离线批处理可能batch=32。

精度校准是INT8量化的核心环节。TensorRT提供熵校准（Entropy Calibration）和最小化均方误差校准（MinMax Calibration）两种算法。校准过程需要准备代表性数据集（通常500-1000张），覆盖真实部署时的数据分布。校准数据不足或分布偏差会导致量化后精度严重下降。

校准数据选择经验：不要使用训练集或测试集，应使用独立采集的、与线上请求分布一致的样本。对于分类模型，确保每个类别都有足够样本；对于检测模型，需包含不同尺度、不同光照条件的图像。校准数据量建议500张起步，若精度不达标可增至2000张。

# trtexec 动态形状 + INT8校准示例
trtexec --onnx=model.onnx \
        --minShapes=input:1x3x224x224 \
        --optShapes=input:8x3x224x224 \
        --maxShapes=input:16x3x224x224 \
        --fp16 \
        --int8 \
        --calib=/data/calibration_images \
        --calibBatchSize=8 \
        --saveEngine=model_fp16_int8.engine

4.3 Engine生成：trtexec与Python API对比

trtexec是TensorRT自带的命令行工具，适合快速验证和基准测试；Python API则适合集成到自动化构建流程中。两者底层调用相同的Builder接口，但Python API提供更细粒度的控制，如自定义校准器、逐层精度设置、动态shape配置等。

对比维度	trtexec	Python API
使用便捷性	命令行参数，无需编程	需要编写Python脚本
动态形状	--minShapes/--optShapes/--maxShapes	profile.set_shape()
INT8校准	--calib指定校准图片目录	自定义Calibrator类
精度控制	全局FP16/INT8	逐层set_precision()
日志输出	详细时间/显存统计	需自行实现Logger
适用场景	快速验证、性能基准	生产环境自动化构建

Python API构建Engine的核心流程如下：

import tensorrt as trt

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)

with open("model.onnx", "rb") as f:
    parser.parse(f.read())

config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)  # 1GB

# 动态形状配置
profile = builder.create_optimization_profile()
profile.set_shape("input", (1,3,224,224), (8,3,224,224), (16,3,224,224))
config.add_optimization_profile(profile)

# 启用FP16
if builder.platform_has_fast_fp16:
    config.set_flag(trt.BuilderFlag.FP16)

# 构建Engine
engine = builder.build_serialized_network(network, config)
with open("model.engine", "wb") as f:
    f.write(engine)

4.4 实战案例：ResNet50从PyTorch到TensorRT的完整转换

某安防项目需要将ResNet50分类模型部署到NVIDIA T4 GPU，要求batch=1时延迟≤2ms，支持动态分辨率（224×224至416×416）。初始PyTorch模型在T4上FP32推理延迟为5.8ms，无法满足需求。

转换步骤与关键数据：

ONNX导出：使用opset=17，动态batch维度，导出后onnxruntime验证余弦相似度0.9998。
trtexec快速验证：FP32 Engine延迟3.2ms，FP16 Engine延迟1.4ms，INT8 Engine延迟0.9ms。FP16已满足需求，决定采用FP16。
动态形状配置：设置min=1×3×224×224，opt=1×3×320×320，max=1×3×416×416。注意opt形状选择320而非最大，因为实际请求中70%为320分辨率。
精度验证：使用1000张真实场景图片对比PyTorch与TensorRT输出，最大误差0.012，Top-5准确率下降0.3%，在可接受范围。
性能基准：最终Engine在T4上平均延迟1.4ms，吞吐量712 FPS（batch=1），满足2ms要求。

风险提示：动态形状范围过大会导致TensorRT为每个可能形状都生成kernel，增加显存占用和构建时间。本案例中max=416×416比opt=320×320多消耗约200MB显存。若线上分辨率稳定，建议缩小max范围以节省显存。另外，INT8虽然延迟更低，但校准数据不足导致准确率下降1.2%，最终放弃INT8方案。

复盘经验：转换过程中最大的坑是torch.onnx.export默认使用training=False，但BatchNorm层在eval模式下行为不同，必须确保模型处于model.eval()状态。此外，动态shape的optShapes选择直接影响性能，应基于线上真实请求分布统计，而非简单取中间值。

4.5 常见误区与优化建议

误区一：认为opset版本越高越好。实际上opset≥17可能引入TensorRT不支持的算子，建议先使用opset=11验证，再逐步升级。
误区二：忽略预处理对齐。PyTorch推理时通常包含Normalize操作，导出ONNX时应将预处理融入模型或确保TensorRT侧执行相同预处理。
误区三：动态shape范围设置过大。过大的范围导致构建时间增长数倍，且可能触发TensorRT的shape优化回退，性能反而下降。
误区四：FP16精度损失盲目接受。对于检测/分割模型，FP16可能导致边界框偏移或mask精度下降，必须逐类评估。
优化建议：使用polygraphy工具进行自动化精度对比和调试，它能快速定位精度下降的层；对于Transformer结构，考虑使用onnx-graphsurgeon进行图融合优化。

ONNX转换与优化是TensorRT部署中最耗时但最关键的环节。掌握动态形状配置、精度校准策略和Engine生成方法，能够显著缩短从训练到上线的周期。建议团队建立标准化的转换流水线，将ONNX导出、验证、Engine构建、精度对比固化为自动化脚本，确保模型迭代时能够快速完成部署验证。