4. PyTorch/TensorFlow模型导出与优化

在 Jetson 上部署模型时，训练框架（PyTorch/TensorFlow）与推理运行时（TensorRT）之间需要一个稳定的中间表示。ONNX 就是这个中间层：它把动态图或静态图固化为带算子集版本的计算图，再由 TensorRT 解析、优化并生成引擎。本节围绕“导出—检查—优化—排错”这条主线，讲清楚每一步实际发生了什么、参数怎么设、结果怎么判断。

一、导出为 ONNX：把训练图固化成可解析的计算图

PyTorch 使用 torch.onnx.export。关键点在于：必须让模型进入 eval() 模式，否则 BatchNorm、Dropout 会以训练行为导出；必须提供与真实推理一致的输入张量（dummy input），因为导出过程会执行一次前向追踪，输入的形状、dtype、设备都会写进图里。

import torch, torchvision

model = torchvision.models.resnet18(weights="IMAGENET1K_V1").eval().cuda()
dummy = torch.randn(1, 3, 224, 224, device="cuda")

torch.onnx.export(
    model, dummy, "resnet18.onnx",
    input_names=["input"], output_names=["logits"],
    opset_version=17,
    do_constant_folding=True,
    dynamic_axes={"input": {0: "batch"}, "logits": {0: "batch"}},
)

这里几个参数直接决定后续能否顺利部署：opset_version 建议 13～17，太低会缺少算子，太高可能超出 TensorRT 解析器支持范围；do_constant_folding=True 让导出器在导出阶段就折叠常量子图；dynamic_axes 声明 batch 维为动态，否则引擎会被锁死在 batch=1。TensorFlow 侧则用 tf2onnx，从 SavedModel 转换：

python -m tf2onnx.convert --saved-model ./saved_model \
  --output model.onnx --opset 17 --inputs input:0 --outputs logits:0

导出后不要急着转 TensorRT，先用 onnx.checker.check_model 验证图结构合法性，再用 onnx.shape_inference.infer_shapes 补全中间张量形状，这两步能提前暴露大量问题。

二、算子兼容性检查：先看 ONNX 支持，再看 TensorRT 支持

兼容性要分两层看。第一层是 PyTorch/TF 算子能否映射到 ONNX 标准算子；第二层是 ONNX 算子能否被 TensorRT 解析器识别。第一层用 onnx.checker 和导出时的告警就能发现；第二层需要对照 TensorRT 的算子支持矩阵，或用 trtexec --onnx=model.onnx --verbose 直接试解析，日志里出现 Unsupported operator 就是明确信号。

常见的不兼容来源包括：自定义 CUDA 算子、torch.nn.functional.interpolate 的某些模式、动态控制流（if/loop）、以及 torchvision.ops.nms 这类后处理算子。处理策略有三种：一是把不支持的部分从图中切出来，用 TensorRT 插件或自定义层实现；二是改写为等价的可支持算子组合；三是把后处理移到 CPU 侧，只把主干网络导出为 ONNX。

三、图优化与常量折叠：导出阶段就能省下的开销

常量折叠（constant folding）指在导出时把输入全为常量的子图直接计算成常量节点。例如 Conv + BatchNorm 在推理模式下 BN 的均值、方差、缩放、偏移都是常量，导出器会把它们融合进卷积权重，图中只剩一个 Conv 节点。这既减少了节点数，也减少了运行时的访存和计算。

除了常量折叠，导出阶段还能做死代码消除（去掉未被输出依赖的节点）、恒等算子消除（如 Identity、Dropout 在推理模式下被移除）。这些优化由 do_constant_folding=True 和导出器自动完成，不需要手工干预。但要注意：如果模型里有依赖运行时数据的动态分支，常量折叠不会生效，此时应检查是否真的需要动态控制流，能否用 torch.where 等可导出算子替代。

四、导出常见错误与修复

错误一：导出后输出与 PyTorch 不一致。最常见原因是模型没有 eval()，BatchNorm 使用了当前 batch 的统计量。修复：导出前显式调用 model.eval()，并确认所有子模块都处于 eval 状态。

错误二：动态 batch 失效。没有设置 dynamic_axes，或设置后 TensorRT 仍按固定 shape 构建。修复：在导出时声明动态维，并在 TensorRT 构建时用 --minShapes/--optShapes/--maxShapes 指定优化配置。

错误三：opset 版本不匹配。导出用了 opset 18，但 Jetson 上的 TensorRT 版本只支持到 opset 17。修复：降低 opset 重新导出，或升级 TensorRT。建议在 Jetson 上先查 dpkg -l | grep tensorrt 确认版本，再决定 opset。

错误四：自定义算子导出失败。PyTorch 自定义算子没有对应的 ONNX 符号。修复：用 torch.onnx.register_custom_op_symbolic 注册映射，或把该算子替换为等价的标准算子组合。

五、一个可复用的导出检查清单

把这张清单固化到导出脚本里，每次导出后自动跑一遍，就能在进入 TensorRT 构建之前把问题定位清楚。下一阶段才是引擎构建与精度校准，而导出的质量直接决定了那一步的顺利程度。