📚 课程目录第4章模型转换

4. 模型转换：将PyTorch模型转换为ONNX，再转换为TensorRT引擎，处理动态形状和精度

4.1 核心原理：为什么需要“PyTorch → ONNX → TensorRT”两级转换

在嵌入式设备上部署模型时，我们很少直接将PyTorch模型交给TensorRT。原因在于PyTorch的动态图机制与TensorRT的静态优化引擎之间存在根本性差异。PyTorch模型在推理时逐算子执行，Python解释器参与调度，这带来了极大的灵活性，但也导致推理延迟高、内存占用大。TensorRT则需要对计算图进行整体分析、层融合、精度校准和内核自动调优，它要求输入一个静态的、完整的、可解析的计算图。

ONNX（Open Neural Network Exchange）恰好充当了中间桥梁。ONNX是一种计算图中间表示，它将PyTorch的动态图“冻结”为静态图，同时保留了算子的语义信息。TensorRT的ONNX解析器（onnx-tensorrt）能够直接读取ONNX文件，将其转换为TensorRT的Network Definition，再经过优化生成引擎。两级转换的流程如下：

  转换链路：PyTorch模型（.pt/.pth）→ ONNX（.onnx）→ TensorRT引擎（.engine/.trt）

这一链路的核心价值在于：PyTorch负责训练和实验，ONNX负责格式统一，TensorRT负责极致优化。每一步都有明确的产出物和验证手段，便于在嵌入式开发中定位问题。

4.2 方法流程：从PyTorch到TensorRT的完整转换步骤

4.2.1 第一步：PyTorch模型导出为ONNX

导出ONNX时，最核心的是动态形状（dynamic shape）的处理。嵌入式设备上，输入图像的尺寸可能因摄像头分辨率、预处理裁剪而改变；batch大小也可能在1和4之间切换。如果导出时固定了形状，后续TensorRT引擎就只能接受该固定尺寸，灵活性大打折扣。

以下代码展示了如何导出支持动态batch和动态宽高的ONNX模型：

import torch
import torch.onnx

# 假设模型已加载并处于eval模式
model = MyModel().eval()
dummy_input = torch.randn(1, 3, 224, 224)  # 示例输入

# 定义动态轴：batch维度=0，高度=2，宽度=3
dynamic_axes = {
    "input": {0: "batch_size", 2: "height", 3: "width"},
    "output": {0: "batch_size"}
}

torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    input_names=["input"],
    output_names=["output"],
    dynamic_axes=dynamic_axes,
    opset_version=17,          # 高版本opset支持更多算子
    do_constant_folding=True,  # 常量折叠优化
    export_params=True
)
print("ONNX导出完成")

经验提示：导出前务必调用 model.eval() 并关闭梯度（torch.no_grad()），否则BatchNorm和Dropout层的行为会不一致，导致导出的ONNX在推理时产生错误结果。

4.2.2 第二步：验证ONNX模型的正确性

导出ONNX后，不能直接跳到TensorRT，必须先验证ONNX的计算结果与原始PyTorch模型是否一致。推荐使用onnxruntime作为参考推理引擎：

import onnxruntime as ort
import numpy as np

# 创建ONNX Runtime会话
sess = ort.InferenceSession("model.onnx", providers=["CPUExecutionProvider"])

# 构造与dummy_input相同的输入
test_input = np.random.randn(1, 3, 224, 224).astype(np.float32)

# 获取输入输出名称
input_name = sess.get_inputs()[0].name
output_name = sess.get_outputs()[0].name

# 运行推理
onnx_output = sess.run([output_name], {input_name: test_input})[0]

# 与PyTorch输出对比
with torch.no_grad():
    torch_output = model(torch.from_numpy(test_input)).numpy()

# 计算最大绝对误差
max_diff = np.max(np.abs(onnx_output - torch_output))
print(f"最大误差: {max_diff:.6f}")
assert max_diff < 1e-4, "ONNX与PyTorch输出差异过大！"

如果误差超过1e-4量级，通常说明ONNX导出时某些算子被错误转换，需要检查模型中的自定义算子或非常规操作。

4.2.3 第三步：ONNX转换为TensorRT引擎

在嵌入式设备上，推荐使用TensorRT的Python API或trtexec命令行工具完成转换。以下代码展示了使用Python API构建引擎的过程，重点处理动态形状和精度设置：

import tensorrt as trt

# 创建logger和builder
logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)
network = builder.create_network(1 << trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
parser = trt.OnnxParser(network, logger)

# 解析ONNX文件
with open("model.onnx", "rb") as f:
    assert parser.parse(f.read()), "ONNX解析失败"

# 配置优化参数
config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)  # 1GB工作空间

# 设置动态形状优化范围
profile = builder.create_optimization_profile()
input_tensor = network.get_input(0)
profile.set_shape(input_tensor.name, 
                  min_shape=(1, 3, 224, 224),   # 最小形状
                  opt_shape=(1, 3, 512, 512),   # 最优形状
                  max_shape=(4, 3, 1024, 1024)) # 最大形状
config.add_optimization_profile(profile)

# 设置精度：FP16
if builder.platform_has_fast_fp16:
    config.set_flag(trt.BuilderFlag.FP16)

# 构建引擎
serialized_engine = builder.build_serialized_network(network, config)
with open("model.engine", "wb") as f:
    f.write(serialized_engine)
print("TensorRT引擎构建完成")

4.3 案例：YOLOv5s在Jetson Orin Nano上的转换实战

以YOLOv5s目标检测模型为例，展示完整的转换流程。YOLOv5s的输入为RGB图像，输出为检测结果，其结构包含CBS模块（Conv+BN+SiLU）、CSP模块和SPPF模块。在转换过程中，最容易出问题的是SiLU激活函数和检测头的解码部分。

实际转换中，我们通常只导出模型的主干+颈部（backbone+neck），而将检测头的解码（anchor解码、NMS）留在TensorRT之外，用自定义CUDA核函数或TensorRT插件实现。这样可以避免ONNX导出时遇到大量非标准算子。

转换阶段	关键操作	常见问题	解决方案
PyTorch → ONNX	设置dynamic_axes，opset≥13	SiLU算子不支持低版本opset	使用opset=17或更高
ONNX验证	onnxruntime对比输出	输出误差>1e-3	检查模型中的Resize/上采样方式
ONNX → TensorRT	设置FP16、动态profile	FP16精度下降导致检测框偏移	对敏感层强制FP32（set_precision）
引擎验证	TensorRT推理对比ONNX	动态形状下输出不一致	确保profile范围覆盖实际输入尺寸

在Jetson Orin Nano上，YOLOv5s转换后的典型性能数据如下：

FP32引擎： 推理延迟约18ms（输入640×640），显存占用约1.2GB
FP16引擎： 推理延迟约7ms，显存占用约0.6GB，mAP下降约0.3%
INT8引擎（需校准）： 推理延迟约4ms，显存占用约0.4GB，mAP下降约1.5%

重要警告：在嵌入式设备上，不要盲目追求INT8量化。如果模型包含大量对精度敏感的层（如检测头、分割头），INT8可能导致输出质量严重下降。建议先使用FP16，确认精度满足需求后再尝试INT8。

4.4 动态形状与精度的深度处理

4.4.1 动态形状的三种模式

TensorRT的优化profile定义了三种形状：min_shape（最小）、opt_shape（最优）和max_shape（最大）。引擎在运行时，会根据实际输入形状选择最接近的优化内核。如果输入形状超出max_shape范围，推理会直接报错。

在嵌入式场景中，建议将opt_shape设置为最常用的推理尺寸。例如，如果摄像头输出为1280×720，经过预处理后输入模型为640×384，那么opt_shape应设为(1, 3, 384, 640)，而不是默认的(1, 3, 224, 224)。

4.4.2 精度处理策略

TensorRT的FP16模式并非将全部层都转为FP16，而是由builder自动决定哪些层适合FP16。对于某些对精度敏感的层（如Softmax、最后的卷积层），可以手动指定为FP32：

# 遍历网络层，对特定层强制FP32
for i in range(network.num_layers):
    layer = network.get_layer(i)
    if layer.name in ["/model.24/m.0/conv", "/model.24/m.1/conv"]:
        layer.precision = trt.float32
        layer.set_output_type(0, trt.float32)

# 同时需要设置config的严格类型标志
config.set_flag(trt.BuilderFlag.OBEY_PRECISION_CONSTRAINTS)

此外，校准数据集对于INT8量化至关重要。校准集应覆盖嵌入式设备实际运行中可能遇到的各种场景（白天/夜晚、不同光照、不同目标密度），通常需要500~1000张代表性图片。校准过程会统计每层激活值的分布，从而确定最优的量化缩放因子。

4.5 常见误区与复盘

误区一：忽略BatchNorm层的融合。在PyTorch中，BatchNorm在训练和推理时的行为不同。导出ONNX前，必须调用model.eval()，否则BatchNorm会使用batch统计量而非全局统计量，导致输出完全错误。更隐蔽的是，如果模型在导出前没有经过torch.jit.script或torch.jit.trace，某些控制流（如if语句）可能无法正确导出。

误区二：动态形状范围设置过大。有些开发者将max_shape设置为(8, 3, 2048, 2048)，希望覆盖所有可能场景。但过大的动态范围会导致TensorRT为每个可能的形状都生成优化内核，显著增加构建时间和显存占用。在Jetson设备上，建议max_shape不超过实际需求的1.5倍。

误区三：FP16精度下降后盲目回退到FP32。FP16推理延迟高、显存占用大，不应轻易放弃。正确的做法是：先定位精度下降严重的层，使用layer.precision = trt.float32单独将这些层保留为FP32，其他层继续使用FP16。这种“混合精度”方案往往能在精度和性能之间取得最佳平衡。

复盘总结：在嵌入式设备上完成模型转换，核心在于“验证每一步”。PyTorch转ONNX后要验证输出一致性；ONNX转TensorRT后要再次验证输出一致性。每一步的误差累积可能导致最终结果完全不可用。建议在开发板上准备一组固定的测试输入（如10张真实图像），在转换的每个阶段都运行一遍，确保输出误差在可接受范围内（通常mAP下降不超过1%）。

最后，将转换好的.engine文件与预处理/后处理代码封装为独立的推理库，供上层应用调用。在嵌入式设备上，引擎文件一旦构建，可以反复加载使用，无需每次启动都重新构建，这能显著缩短应用的启动时间。