3. ONNX模型导出与优化：从PyTorch/TensorFlow导出ONNX，使用onnx-simplifier简化模型，Fix动态轴

好，咱们进入第三章。这一章聊的是ONNX导出和优化。

说实话，ONNX这东西，在TensorRT的整个工作流里，就是个中间桥梁。桥要是搭歪了，后面跑起来全是坑。我见过太多人，模型在PyTorch里跑得飞起，一转到TensorRT就各种报错，十有八九是ONNX导出这步没做好。

3.1 从PyTorch导出ONNX

PyTorch导出ONNX，官方提供了torch.onnx.export这个接口。看起来就一行代码，但里面的门道不少。

先看一个最基础的例子：

import torch
import torch.onnx

# 假设我们有一个训练好的模型
model = torch.load('resnet50.pth')
model.eval()

# 创建一个 dummy input
dummy_input = torch.randn(1, 3, 224, 224)

# 导出 ONNX
torch.onnx.export(
    model,               # 模型
    dummy_input,         # 输入张量
    'resnet50.onnx',     # 输出文件名
    export_params=True,  # 是否导出参数
    opset_version=11,    # ONNX算子集版本
    do_constant_folding=True,  # 常量折叠优化
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={
        'input': {0: 'batch_size'},
        'output': {0: 'batch_size'}
    }
)

这里有几个关键参数，我重点说一下：

• opset_version：我个人习惯用11或13。太低了有些算子不支持，太高了有些老版本的TensorRT不兼容。11是个比较稳妥的选择。
• do_constant_folding：这个建议打开。它会把一些计算图里可以提前算好的常量折叠掉，减少运行时计算量。
• dynamic_axes：这个就是咱们要说的动态轴。后面会详细讲。

3.2 从TensorFlow导出ONNX

TensorFlow导出ONNX稍微麻烦一点。需要用tf2onnx这个工具。

安装很简单：

pip install tf2onnx

然后命令行导出：

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

或者用Python API：

import tf2onnx
import tensorflow as tf

model = tf.saved_model.load('./saved_model')
spec = (tf.TensorSpec((None, 224, 224, 3), tf.float32, name="input"),)
output_path = "./model.onnx"

model_proto, _ = tf2onnx.convert.from_keras(
    model, 
    input_signature=spec, 
    opset=11,
    output_path=output_path
)

嗯，这里要注意一点。TensorFlow的模型结构跟PyTorch不太一样，它的输入输出命名方式比较特殊。我建议导出后先用Netron看一眼，确认输入输出名字对不对。

3.3 使用onnx-simplifier简化模型

导出的ONNX模型，往往有很多冗余操作。比如一些恒等映射、多余的reshape、不必要的cast等等。这些操作在推理时不会带来任何好处，反而会增加计算开销。

这时候就需要onnx-simplifier出场了。

安装：

pip install onnx-simplifier

使用：

python -m onnxsim model.onnx model_simplified.onnx

或者用Python：

import onnx
from onnxsim import simplify

model = onnx.load('model.onnx')
model_simplified, check = simplify(model)

if check:
    onnx.save(model_simplified, 'model_simplified.onnx')
    print("简化成功！")
else:
    print("简化失败，请检查模型结构")

简化后的模型，节点数通常会减少30%-50%。我优化过一个BERT模型，从1200个节点简化到了700多个，推理速度提升了将近20%。

3.4 Fix动态轴

动态轴，说白了就是让模型支持可变大小的输入。比如你的模型要处理不同batch size的数据，或者不同尺寸的图片。

在ONNX里，动态轴是通过dynamic_axes参数来指定的。看前面的例子：

dynamic_axes={
    'input': {0: 'batch_size'},
    'output': {0: 'batch_size'}
}

这里的意思是，input的第0维（batch维度）是动态的，名字叫batch_size。output也一样。

但动态轴有个坑——不是所有算子都支持动态形状。比如Reshape、Gather这些算子，如果输入形状是动态的，可能会导出失败。

我遇到过最典型的一个问题：模型里有个Flatten操作，输入是动态batch size，结果导出时报错说Flatten不支持动态形状。后来我把Flatten换成了Reshape，手动指定了形状，才搞定。

解决方案是这样的：

# 原来的写法
x = torch.flatten(x, start_dim=1)

# 改成手动reshape
batch_size = x.shape[0]
x = x.view(batch_size, -1)

这样导出的ONNX就能正确处理动态batch了。

3.5 导出后的验证

模型导出后，一定要验证一下。我习惯做三件事：

1. 用Netron查看计算图：确认输入输出节点、算子类型是否正确。
2. 用onnxruntime跑一下：对比PyTorch/TensorFlow的输出，看误差是否在可接受范围内。
3. 检查动态轴：用不同形状的输入测试，确保动态轴正常工作。

验证代码示例：

import onnxruntime
import numpy as np

# 加载ONNX模型
session = onnxruntime.InferenceSession('model_simplified.onnx')

# 准备输入
input_name = session.get_inputs()[0].name
input_data = np.random.randn(2, 3, 224, 224).astype(np.float32)

# 推理
outputs = session.run(None, {input_name: input_data})

# 对比结果
# 这里假设你已经有PyTorch的输出结果
print("ONNX输出:", outputs[0][0][:5])
print("PyTorch输出:", torch_output[0][:5])

如果误差在1e-5以内，基本就没问题了。

3.6 避坑指南

最后，分享几个我踩过的坑：

• 算子兼容性：有些PyTorch算子（比如torch.einsum）在ONNX里没有直接对应，导出时会自动拆成多个算子，可能导致性能下降。我建议尽量用ONNX支持的算子。
• 动态轴命名：动态轴的名字不要用中文或特殊字符，否则某些工具解析会出错。
• opset版本：不同版本的opset支持的算子不同。如果导出失败，试试降低opset版本。

嗯，这一章的内容就这些。下一章咱们聊聊TensorRT的模型解析和构建，那才是真正开始加速的地方。