4、ONNX模型导出：从PyTorch/TensorFlow导出ONNX，常见踩坑点

说实话，ONNX导出这事儿，看着简单，做起来全是坑。我刚开始做部署那会儿，觉得不就是个模型格式转换嘛，结果一跑推理，输出全是NaN，查了整整两天。嗯，今天就把我这些年踩过的坑，一个一个给你掰扯清楚。

4.1 动态轴与静态轴的陷阱

ONNX导出时，最容易被忽略的就是动态轴的问题。PyTorch模型默认是动态的，但ONNX默认是静态的。你想想看，训练时batch size可以随便设，但导出ONNX时如果不指定，它就固定死了。

正确的做法是这样：

import torch
import torch.onnx

# 假设你的模型输入是 [batch, 3, 224, 224]
dummy_input = torch.randn(1, 3, 224, 224)

# 关键：指定动态轴
dynamic_axes = {
    'input': {0: 'batch_size'},   # 第0维是动态的
    'output': {0: 'batch_size'}
}

torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    input_names=['input'],
    output_names=['output'],
    dynamic_axes=dynamic_axes,
    opset_version=11  # 推荐用11或更高
)

我个人习惯把batch size、height、width都设成动态的。但要注意，动态轴越多，TensorRT优化效果越差。所以，能固定就固定，别贪心。

4.2 算子兼容性问题

这是个大坑。PyTorch里很多高级算子，ONNX根本不支持。比如 torch.einsum、torch.where 的某些用法，导出时直接报错。

我在项目中遇到过最典型的一个：torch.nn.functional.grid_sample。这个算子做图像变形时特别好用，但ONNX不支持。怎么办？

• 方案一：用 torch.onnx.export 的 custom_opsets 参数注册自定义算子
• 方案二：改写模型，用ONNX支持的算子替代（比如用 torch.nn.functional.interpolate + 仿射变换）
• 方案三：导出时用 torch.onnx.export 的 operator_export_type 参数，把不支持的算子拆成基本运算

4.3 常量折叠与输入输出命名

ONNX导出时，PyTorch会自动做常量折叠。什么意思？就是会把一些固定计算提前算好，写进模型里。这本来是好事，但有时候会出问题。

举个例子：如果你的模型里有个 if 分支，依赖某个输入值，常量折叠可能会把分支给优化掉。嗯，这会导致推理结果完全不对。

解决办法：用 torch.onnx.export 的 do_constant_folding=False 参数关掉它。

torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    do_constant_folding=False,  # 关掉常量折叠
    ...
)

另外，输入输出命名也很重要。我建议用有意义的名称，比如 'input'、'output'，别用默认的 '0'、'1'。因为TensorRT解析时，名字对不上会直接报错。

4.4 TensorFlow导出ONNX的特殊问题

TensorFlow导出ONNX，我个人觉得比PyTorch更麻烦。主要问题有两个：

问题	原因	解决办法
图结构复杂	TF的静态图包含大量控制流节点	用 tf2onnx 工具，设置 --opset 13
输入形状不确定	TF模型输入经常是 None	导出时显式指定 input_signature
自定义层无法导出	TF的自定义层ONNX不认识	用 tf.function 包装，或手动实现ONNX算子

我曾经有个项目，TF模型里用了 tf.while_loop，导出ONNX后，TensorRT直接报 Unsupported operator。最后没办法，只能把循环展开成固定次数的 tf.map_fn。

4.5 精度问题：FP16与INT8导出

导出ONNX时，默认是FP32。但如果你打算用TensorRT做FP16或INT8推理，导出时就要注意了。

我建议：先导出FP32的ONNX，再用TensorRT做精度校准。别直接在导出时转FP16，因为ONNX的FP16支持不完善，容易出精度损失。

4.6 避坑指南：我踩过的那些坑

最后，分享几个我亲身经历过的坑，你遇到了可以直接抄答案：

• 坑1：导出时 opset_version 太低。我建议至少用 opset_version=11，因为很多新算子需要高版本支持。
• 坑2：忘记设置 input_names 和 output_names。TensorRT解析时，默认名字是 'input.1'、'output.1'，容易搞混。
• 坑3：模型里有 torch.jit.script 或 @torch.jit.script 装饰器。这些脚本化的函数，导出ONNX时经常出问题，建议先去掉。
• 坑4：TensorFlow的 tf.keras.layers.BatchNormalization 导出后，推理结果和训练时不一致。这是因为ONNX的BN实现和TF有差异，建议用 tf.nn.batch_normalization 替代。

嗯，差不多就这些。ONNX导出这事儿，说白了就是「试错-修正-再试」的过程。你只要记住：先导出FP32，验证精度，再考虑优化。别一上来就搞FP16，否则出了问题你都不知道是模型的问题还是导出的问题。

下次咱们聊聊TensorRT的模型解析和网络构建，那又是另一番天地了。