3、模型转换基础：ONNX作为中间表示、从PyTorch导出ONNX的常见陷阱、使用trtexec进行快速转换
模型转换，说白了就是让模型换个语言说话。PyTorch讲的是Python，TensorRT讲的是CUDA。中间得有个翻译官——ONNX就是这个角色。
我个人习惯把ONNX看作模型的"通用语言"。不管你是PyTorch、TensorFlow还是其他框架训练出来的模型，只要能转成ONNX，就能被TensorRT吃进去。嗯，这里要注意，ONNX不是万能的，但它是目前最靠谱的中间格式。
3.1 ONNX：为什么选它做中间人？
ONNX全称是Open Neural Network Exchange，微软和Meta一起搞的。它的核心价值就一个：解耦。训练框架和推理框架不再绑死。
我在项目中遇到过好几次这种情况：客户用PyTorch训练了一个检测模型，但部署环境只支持TensorRT。没有ONNX的话，你得重写一遍模型结构，那叫一个痛苦。有了ONNX，导出一次，到处跑。
ONNX的优势很明显：
- 算子标准化：定义了统一的算子集，比如Conv、Relu、BatchNormalization
- 静态图结构：计算图是固定的，方便优化和剪枝
- 类型系统：每个张量的数据类型、形状都明确标注
- 工具链成熟：从导出到可视化，再到优化，生态很完善
核心要点：ONNX不是用来训练的，它是用来"运输"模型的。你训练时用PyTorch的灵活动态图，部署时用ONNX的静态图，各取所长。
3.2 从PyTorch导出ONNX：看着简单，坑不少
PyTorch导出ONNX就一行代码：torch.onnx.export()。但这一行背后，藏着无数个坑。我刚开始做部署时，光导出就折腾了两天。
先看一个标准的导出流程：
import torch
import torchvision.models as models
# 加载模型
model = models.resnet18(pretrained=True)
model.eval()
# 构造一个示例输入
dummy_input = torch.randn(1, 3, 224, 224)
# 导出ONNX
torch.onnx.export(
model, # 模型
dummy_input, # 示例输入
"resnet18.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'}
}
)
代码看着挺规整，对吧？但实际跑起来，问题一个接一个。
3.3 常见陷阱：我踩过的那些坑
陷阱一：动态shape处理不当
PyTorch模型天生支持动态shape，但ONNX是静态图。如果你不指定dynamic_axes，导出的ONNX就固定死了batch size=1。部署时换个batch size，直接报错。
我曾经有个项目，模型在测试集上跑得好好的，一上线就崩。查了半天，发现是ONNX里batch size写死了。后来加上dynamic_axes才解决。
警告：动态shape会降低TensorRT的优化效果。能固定尽量固定，实在需要动态，只对batch维度开动态，其他维度尽量固定。
陷阱二：算子不支持
PyTorch有些算子ONNX不支持，比如torch.einsum、某些自定义的CUDA扩展。导出时会报错：Unsupported operator。
我建议导出前先检查模型里有没有冷门算子。常用的Conv、Relu、Pooling都没问题，但像torch.topk、torch.meshgrid这些，不同opset版本支持情况不一样。
解决办法：
- 升级opset版本（一般用11或以上）
- 用等价算子替换（比如把einsum拆成matmul和transpose）
- 实在不行，写自定义ONNX算子（但很麻烦，不推荐）
陷阱三：控制流被固化
PyTorch模型里如果有if、for循环，导出ONNX时会被展开成静态图。什么意思？就是循环次数固定死了。
我记得有个NLP模型，里面有个while循环处理变长序列。导出后，所有序列都被填充到最大长度，推理速度直接掉了30%。
解决方案：用torch.onnx.export的dynamic_axes配合torch.nn.utils.rnn.pad_packed_sequence来处理变长输入。
陷阱四：训练和推理模式混用
导出前一定要调用model.eval()。我见过有人忘了这步，结果导出的ONNX里包含了Dropout和BatchNorm的训练行为。推理时结果完全不对。
小技巧：导出前用torch.no_grad()包裹，确保没有梯度计算。同时用model.eval()关闭训练特有的层。
3.4 使用trtexec快速转换
ONNX导出来了，怎么转成TensorRT的engine？最直接的方式就是用trtexec。这是TensorRT自带的一个命令行工具，功能很全。
基本用法：
trtexec --onnx=resnet18.onnx \
--saveEngine=resnet18.engine \
--workspace=1024 \
--fp16
参数说明：
| 参数 | 作用 | 我的建议 |
|---|---|---|
| --onnx | 指定ONNX模型路径 | 路径不要有中文和空格 |
| --saveEngine | 输出engine文件路径 | 命名带上精度和batch size |
| --workspace | 工作空间大小（MB） | 大模型给2048以上 |
| --fp16 | 启用FP16推理 | 精度损失小，速度提升明显 |
| --int8 | 启用INT8推理 | 需要校准数据集，精度可能下降 |
| --best | 自动选择最优精度 | 适合快速测试 |
trtexec不只是转换工具，它还能做性能测试：
trtexec --loadEngine=resnet18.engine \
--shapes=input:1x3x224x224 \
--iterations=100 \
--warmUp=10
这个命令会加载engine，跑100次推理，输出平均延迟、吞吐量等指标。我每次转换完都会跑一遍，看看性能是否达标。
实战经验：trtexec的--verbose参数很有用。加上它，你能看到每一层的转换细节，包括哪些层用了FP16、哪些层回退到了FP32。这对性能分析帮助很大。
3.5 转换流程总结
整个流程走下来，其实就三步：
- 导出ONNX：PyTorch模型转成ONNX，注意动态shape和算子兼容性
- 验证ONNX：用
onnxruntime跑一遍，确保结果和PyTorch一致 - 转成engine：用trtexec或TensorRT Python API转换，选择合适精度
我建议每步都做验证。导出后先跑ONNX Runtime，再转engine，最后对比精度。一步出错，后面全白干。
嗯，模型转换这块，说白了就是细心活。坑不少，但踩过一次就记住了。下一章我们会深入TensorRT的优化策略，看看怎么把engine的性能榨到极致。