第四章：ONNX与模型解析

好，咱们进入第四章。这一章我打算聊聊ONNX和模型解析。说实话，很多做TensorRT部署的朋友，最后卡住的地方往往不是推理优化本身，而是模型根本解析不过去。嗯，ONNX就是那个中间桥梁，桥要是塌了，后面再好的优化也白搭。

4.1 ONNX基础：你总得知道它长什么样

ONNX，全称Open Neural Network Exchange。说白了，它就是AI模型界的“通用语言”。你想想看，PyTorch训练出来的模型，TensorRT不认识；TensorFlow的模型，ONNX Runtime也不一定直接吃。那怎么办？大家约定一个中间格式，ONNX就是这个格式。

我个人习惯把ONNX理解成一个“计算图描述文件”。它里面记录了：

  • 计算节点（Node）：比如Conv、Relu、Add这些算子
  • 张量（Tensor）：节点之间的数据流动
  • 图结构（Graph）：节点怎么连在一起的
  • 元数据（Metadata）：模型版本、作者、opset版本等

这里有个小细节，opset版本特别重要。我在项目中遇到过好几次，PyTorch导出的ONNX用了opset 17，但TensorRT只支持到opset 13，结果解析直接报错。所以导出时我一般会指定opset版本，比如：

torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    opset_version=13,  # 我习惯用13，兼容性好
    input_names=["input"],
    output_names=["output"]
)
小提示：如果你不确定用哪个opset版本，去查一下当前TensorRT版本支持的ONNX opset范围。别盲目追新。

4.2 使用ONNX Runtime验证模型

模型导出成ONNX之后，第一件事不是扔给TensorRT，而是先用ONNX Runtime跑一遍。为什么？因为你要确认导出的ONNX本身没问题。我曾经有一次，模型导出时dynamic axes设置错了，结果TensorRT解析一直报shape mismatch，查了两天才发现是ONNX本身就有问题。

验证代码其实很简单：

import onnxruntime as ort
import numpy as np

# 创建推理会话
sess = ort.InferenceSession("model.onnx")

# 获取输入输出信息
input_name = sess.get_inputs()[0].name
output_name = sess.get_outputs()[0].name

# 构造假数据
dummy_input = np.random.randn(1, 3, 224, 224).astype(np.float32)

# 跑一次推理
output = sess.run([output_name], {input_name: dummy_input})

print("ONNX Runtime推理成功，输出shape:", output[0].shape)

这里我特别强调一点：输入数据的dtype一定要和ONNX模型要求的一致。很多人在这一步翻车，明明模型要求float32，你塞了个float64进去，ONNX Runtime不会报错，但TensorRT会。

核心原则：ONNX Runtime验证通过，是TensorRT解析成功的前提。这一步别省。

4.3 TensorRT解析ONNX的流程

好，现在ONNX验证通过了，咱们来看看TensorRT是怎么吃下这个ONNX的。流程其实不复杂，我拆成几步讲：

  1. 读取ONNX文件：TensorRT调用ONNX解析器（onnxparser）读取.proto文件
  2. 构建计算图：解析器把ONNX的节点映射成TensorRT的ILayer
  3. 优化与编译：TensorRT对图进行融合、精度校准、内存优化
  4. 生成引擎：最终生成一个序列化的TensorRT引擎文件（.trt或.engine）

代码实现大概是这样的：

import tensorrt as trt

# 创建logger和builder
logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)

# 创建网络定义
network = builder.create_network(
    1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
)

# 创建ONNX解析器
parser = trt.OnnxParser(network, logger)

# 解析ONNX文件
with open("model.onnx", "rb") as f:
    if not parser.parse(f.read()):
        for error in range(parser.num_errors):
            print(parser.get_error(error))
        raise RuntimeError("ONNX解析失败")

# 构建引擎
config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)  # 1GB
engine = builder.build_serialized_network(network, config)
注意：EXPLICIT_BATCH标志必须加。不加的话，TensorRT会默认使用implicit batch模式，很多动态shape操作会报错。我刚开始做的时候就被这个坑过。

4.4 TensorRT解析ONNX的限制

说到限制，这才是本章的重点。TensorRT不是万能的，它解析ONNX时有几个硬伤：

限制类型 具体表现 我的建议
算子支持 某些ONNX算子TensorRT不支持 用自定义插件（后面章节会讲）
动态shape 部分动态操作解析失败 尽量固定shape，或用优化profile
控制流 If、Loop等节点解析困难 导出时禁用控制流，或用onnx-simplifier
数据类型 不支持double、int64等类型 导出时强制转float32

我举个例子。有一次我导出一个BERT模型，里面有个Gather操作，ONNX Runtime跑得好好的，但TensorRT解析时报错："Unsupported ONNX node type: Gather"。后来查了半天，发现是Gather的indices是int64类型，TensorRT只支持int32。解决办法就是在导出前把indices转成int32。

避坑指南：我曾经用onnx-simplifier工具简化模型，它能自动处理很多不兼容问题。命令很简单：python -m onnxsim model.onnx model_sim.onnx。建议每次导出后都跑一遍。

4.5 实战经验总结

最后，我分享几个实战中总结的经验：

  • 先跑ONNX Runtime：这是最低成本的验证方式
  • 用Netron可视化：导出ONNX后，用Netron打开看看图结构，一眼就能发现异常节点
  • 关注opset版本：TensorRT 8.x一般支持opset 13-16，太新或太旧都可能出问题
  • 不要迷信一键转换：很多模型需要手动调整才能被TensorRT正确解析

嗯，这一章就到这里。ONNX解析是TensorRT开发的第一道坎，跨过去之后，后面的优化工作才会顺利。下一章咱们聊聊自定义插件，那才是真正展示技术的地方。

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