3、ONNX基础：ONNX格式介绍、PyTorch模型导出为ONNX、ONNX模型可视化

好，咱们进入第三章。这一章聊的是ONNX，一个在模型部署里绕不开的中间格式。

说实话，我刚开始做部署那会儿，对ONNX的理解就是「把PyTorch模型转一下就能跑TensorRT了」。后来踩了不少坑才明白，ONNX本身就是一个精密的中间表示，它不仅仅是格式转换，更是模型优化的第一步。

3.1 ONNX格式介绍

ONNX全称是Open Neural Network Exchange，说白了就是给不同深度学习框架之间搭的一座桥。你想想看，PyTorch训练好的模型，想跑到TensorRT上，中间总得有个大家都认的格式吧？ONNX就是这个角色。

我个人习惯把ONNX理解成「模型的通用语言」。不管你是用PyTorch、TensorFlow还是Keras，最终都能导出成ONNX。然后下游的推理引擎，比如TensorRT、OpenVINO、ONNX Runtime，都能直接吃进去。

ONNX的核心结构其实不复杂，主要就两块：

  • 计算图（Graph）：描述模型的计算流程，包括节点（算子）和边（张量流动）
  • 元数据（Metadata）：模型的版本、作者、输入输出信息等

每个节点代表一个算子，比如Conv、Relu、Add这些。边就是张量数据。整个图是静态的，也就是说，一旦导出，计算路径就固定了。这一点和PyTorch的动态图不一样——嗯，这里要注意，动态图里的控制流（比如if-else、循环）在导出ONNX时可能会出问题。

关键点：ONNX是静态图，不支持动态控制流。如果你的模型里有条件分支，导出时需要用torch.onnx.export的dynamic_axes参数来处理，或者干脆把控制流逻辑放到模型外部。

我在项目中遇到过一个问题：一个用PyTorch写的检测模型，里面有个根据输入尺寸动态调整anchor的逻辑。导出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'}
    }
)

print("ONNX导出完成！")

这段代码看着简单，但有几个坑我得提醒你：

  • model.eval()一定要加：不切到eval模式，BatchNorm和Dropout的行为会不一样，导出的ONNX推理结果可能和训练时对不上。我曾经因为这个排查了一整天，最后发现是忘了加eval()。
  • dummy_input的shape要和实际输入一致：不仅仅是维度数，每个维度的含义也要对。比如图像输入是NCHW格式，别搞成NHWC。
  • opset_version别选太高：我一般用11或13。版本太高，有些推理引擎可能不支持。版本太低，有些新算子又用不了。折中一下，11是个稳妥的选择。
  • dynamic_axes处理动态batch：如果你的模型需要支持不同batch size的输入，一定要设置dynamic_axes。否则导出的ONNX会固定batch size为1。

小技巧：导出后可以用onnx.checker.check_model验证一下ONNX模型是否合法。我每次导出完都会跑一遍这个检查，能提前发现不少问题。

import onnx

# 加载并检查ONNX模型
onnx_model = onnx.load("resnet18.onnx")
onnx.checker.check_model(onnx_model)
print("ONNX模型检查通过！")

说到验证，我记得有一次导出一个BERT模型，checker报了个「Graph must be acyclic」的错误。查了半天，发现是模型里有个循环依赖。后来用onnx.shape_inference.infer_shapes做了形状推断，才定位到问题。

3.3 ONNX模型可视化

模型导出来了，怎么确认它长什么样？可视化是个好办法。我个人习惯用Netron，一个轻量级的ONNX可视化工具。

Netron的使用方式很简单：

  • 网页版：直接打开 netron.app，把.onnx文件拖进去
  • 本地版：pip install netron，然后 netron resnet18.onnx 启动本地服务

可视化之后你能看到什么？

  • 整个计算图的拓扑结构，每个节点是什么算子
  • 每个张量的shape和数据类型
  • 输入输出的名称和维度信息
  • 节点的属性参数，比如卷积的kernel size、stride等

我一般用可视化做三件事：

  1. 检查结构是否正确：看看模型是不是按预期搭的，有没有多余的节点
  2. 定位问题节点：如果TensorRT跑不了，可视化能帮你快速找到不支持的那个算子
  3. 优化参考：看看哪些地方可以合并，比如Conv+Bn+Relu这种模式

注意：Netron只能看结构，不能看权重值。如果你想检查权重，得用onnx.numpy_helper.to_array把权重提取出来看。

除了Netron，还有一个命令行工具onnx_graphsurgeon也可以用来分析ONNX图。不过说实话，日常调试用Netron就够了，图形化界面直观得多。

好，这一章的内容就这些。ONNX是模型部署的基石，把这块搞扎实了，后面转TensorRT才能顺风顺水。下一章咱们就正式进入TensorRT的世界，聊聊怎么把ONNX模型跑起来。

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