2、ONNX导出实战：从PyTorch导出ONNX、从TensorFlow导出ONNX、从Keras导出ONNX、常见导出问题与修复

好，咱们直接进入正题。

ONNX 这东西，说白了就是模型界的“通用语言”。你拿 PyTorch 训好的模型，想放到 TensorRT 上跑？中间必须过一道 ONNX。我见过太多人卡在这一步，模型训得挺好，一导出就报错，心态直接崩了。

今天我就把三种主流框架的导出方法掰开揉碎讲清楚。顺便把我踩过的坑也抖出来，你遇到类似问题直接抄作业就行。

2.1 从PyTorch导出ONNX

PyTorch 导出 ONNX 是最常见的场景。官方提供了 torch.onnx.export 接口，用起来不算复杂，但细节决定成败。

2.1.1 基本导出流程

先看一个标准示例。假设你有一个训练好的 ResNet-18 分类模型：

import torch
import torchvision.models as models

# 加载模型，设为评估模式
model = models.resnet18(pretrained=True)
model.eval()

# 构造一个 dummy 输入
dummy_input = torch.randn(1, 3, 224, 224)

# 导出 ONNX
torch.onnx.export(
    model,                     # 模型
    dummy_input,               # 示例输入
    "resnet18.onnx",           # 输出路径
    export_params=True,        # 导出参数
    opset_version=11,          # opset 版本
    do_constant_folding=True,  # 常量折叠优化
    input_names=["input"],     # 输入名称
    output_names=["output"],   # 输出名称
    dynamic_axes={
        "input": {0: "batch_size"},
        "output": {0: "batch_size"}
    }
)

这里我重点说几个参数。

opset_version：我个人习惯用 11 或 13。opset 版本越高，支持的操作越多，但兼容性可能下降。TensorRT 8.x 对 opset 11 支持最稳，我一般选这个。

dynamic_axes：这个很关键。如果你想让模型支持动态 batch 大小，必须配这个参数。我在项目中遇到过客户要求 batch size 从 1 到 32 动态变化，没配 dynamic_axes 的话，导出后只能固定 batch，那叫一个尴尬。

小技巧：导出前务必调用 model.eval()。我见过有人忘了这步，结果导出的 ONNX 里包含了 dropout 和 batchnorm 的训练行为，推理结果完全不对。

2.1.2 处理动态输入尺寸

除了动态 batch，有时候输入图片的宽高也是变化的。比如目标检测模型，输入可能是 640x640，也可能是 1280x720。

这时候 dynamic_axes 要这样配：

dynamic_axes = {
    "input": {0: "batch_size", 2: "height", 3: "width"},
    "output": {0: "batch_size"}
}

嗯，这里要注意：不是所有算子都支持动态尺寸。像 Resize、GridSample 这些操作，在动态尺寸下容易出问题。我建议你先用固定尺寸导出，验证通过后再改动态。

2.2 从TensorFlow导出ONNX

TensorFlow 导出 ONNX 稍微绕一点。因为 TF 的模型格式有好几种：SavedModel、H5、Checkpoint、Frozen Graph。我推荐用 SavedModel 格式作为起点。

2.2.1 使用 tf2onnx 工具

社区最常用的工具是 tf2onnx。安装很简单：

pip install tf2onnx

导出命令如下：

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

这里有个坑：TF 的 tensor 命名规则是 操作名:输出索引。比如 input:0 表示第一个输入 tensor。你最好用 saved_model_cli show 命令查看模型的输入输出名称：

saved_model_cli show --dir ./saved_model --tag_set serve --signature_def serving_default

我曾经帮一个团队排查问题，他们导出的 ONNX 推理结果全是 NaN。查了半天，发现是输入名称写错了，导出的 ONNX 根本没连上正确的输入节点。这种低级错误，排查起来最费时间。

2.2.2 处理 TF 特有操作

TF 有一些自己的算子，比如 FusedBatchNorm、Identity。这些在 ONNX 里没有直接对应。tf2onnx 会自动做转换，但有时候会失败。

我遇到过最典型的问题是 tf.while_loop 导不出来。ONNX 不支持循环控制流（除非你用 opset 14+ 的 Loop 算子）。解决方案是：在导出前把循环展开，或者用 tf.map_fn 替代。

警告：TensorFlow 2.x 的 Keras 模型导出时，记得设置 model.run_eagerly = False。否则导出的 ONNX 会包含 eager 模式的调试信息，体积大且效率低。

2.3 从Keras导出ONNX

Keras 导出 ONNX 有两种路径：一是通过 TF 后端走 tf2onnx，二是用 keras2onnx 工具。我推荐前者，因为 keras2onnx 已经很久没更新了。

2.3.1 通过 tf2onnx 导出 Keras 模型

假设你有一个训练好的 Keras 模型：

import tensorflow as tf
import tf2onnx

model = tf.keras.applications.MobileNetV2(weights="imagenet")
model.run_eagerly = False  # 重要！

# 指定输入签名
spec = (tf.TensorSpec((None, 224, 224, 3), tf.float32, name="input"),)

# 导出
output_path = "mobilenetv2.onnx"
model_proto, _ = tf2onnx.convert.from_keras(model, input_signature=spec, opset=13)
with open(output_path, "wb") as f:
    f.write(model_proto.SerializeToString())

这里 input_signature 必须指定。我见过有人不写这个，结果导出的 ONNX 输入形状是 unknown，TensorRT 编译直接报错。

2.3.2 Keras 特有的 Lambda 层问题

Keras 里经常用 Lambda 层封装自定义操作。比如：

model.add(Lambda(lambda x: x * 2.0))

这种层导出时，tf2onnx 会尝试把 lambda 函数转换成 ONNX 算子。但如果 lambda 里用了 Python 的 if 或 for，转换就会失败。

我的建议是：尽量避免在 Keras 模型里用 Lambda 层。实在要用，就用 tf.keras.layers.Layer 子类化，把逻辑写在 call 方法里。这样导出更稳定。

2.4 常见导出问题与修复

这部分是我最想分享的。导出 ONNX 时遇到的问题，翻来覆去就那么几类。我把它们总结成了一张表：

问题现象 可能原因 解决方案
导出报错：Unsupported operator 模型使用了 ONNX 不支持的算子 升级 opset 版本，或替换为等价操作
推理结果与原始模型不一致 模型未设为 eval 模式，或存在随机操作 调用 model.eval()，设置随机种子
ONNX 文件体积过大 包含了大量常量或未折叠的节点 开启 do_constant_folding，或使用 onnx-simplifier
动态 batch 导出失败 某些算子不支持动态形状 用固定 batch 导出，或替换算子
TensorRT 编译报 shape 错误 ONNX 中存在动态维度未指定 在 dynamic_axes 中明确所有动态维度

2.4.1 算子不支持怎么办？

这是最头疼的问题。比如 PyTorch 里的 torch.nn.functional.grid_sample，在 ONNX opset 11 里没有对应算子。

我的处理思路是：

  1. 升级 opset：先试试 opset 13 或 16，新版本往往增加了算子支持。
  2. 替换实现：比如 grid_sample 可以用 torch.nn.functional.interpolate + 仿射变换替代。
  3. 自定义算子：实在不行就写一个 ONNX 自定义算子，但 TensorRT 那边也得注册，工作量比较大。

我记得有一次，一个客户用了 torch.topk，ONNX 导出没问题，但 TensorRT 编译报错。后来发现是 opset 版本太低，升级到 13 就解决了。

2.4.2 精度对不上怎么办？

导出后推理结果和原始模型差很多？别慌，先检查这几步：

  • 确认模型处于 eval 模式
  • 确认输入数据预处理一致（归一化、通道顺序等）
  • 用 onnxruntime 跑一遍 ONNX，看结果是否一致
  • 如果 onnxruntime 结果正确但 TensorRT 不对，那问题出在 TensorRT 编译阶段

我曾经遇到一个 case：PyTorch 模型用了 torch.nn.Upsample，导出 ONNX 后，onnxruntime 结果正常，但 TensorRT 编译后结果偏了。排查发现是 TensorRT 对 Upsample 的 align_corners 参数处理有 bug。解决方案是把 Upsample 换成 torch.nn.functional.interpolate 并指定 mode='nearest'。

核心原则：导出 ONNX 后，一定要用 onnxruntime 做一次验证。这一步能帮你区分问题是出在 ONNX 导出阶段，还是 TensorRT 编译阶段。

2.4.3 使用 onnx-simplifier 优化

导出的 ONNX 文件里经常有很多冗余节点，比如 Identity、Cast、Shape 等。这些节点不影响结果，但会增加 TensorRT 编译时间，甚至导致编译失败。

我推荐用 onnx-simplifier 做一次简化：

pip install onnx-simplifier

python -m onnxsim model.onnx model_simplified.onnx

这个工具会做常量折叠、节点合并、死代码消除。我处理过的模型，简化后体积能减少 30%~50%，编译时间也明显缩短。

嗯，到这里，ONNX 导出的核心内容就差不多了。你可能会问：为什么我导出的 ONNX 在 TensorRT 里还是跑不起来？别急，下一章我会专门讲 TensorRT 的编译流程和常见报错处理。到时候你会发现，很多问题其实在导出阶段就已经埋下了伏笔。

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