4、ONNX中间表示：ONNX的诞生背景、算子集定义、模型结构解析、如何将PyTorch/TF模型导出为ONNX。
各位同学，咱们今天聊聊ONNX。说实话，这玩意儿在AI芯片部署的链路里，是个绕不开的坎儿。你想想看，算法团队用PyTorch，训练团队用TensorFlow，到了我们芯片这边，总不能每种框架都去适配一遍吧？那不得累死。ONNX就是来解决这个问题的。
4.1 ONNX的诞生背景：为什么我们需要一个“中间人”
早些年，AI框架那是群雄割据。PyTorch、TensorFlow、Caffe、MXNet……每个框架都有自己的模型存储格式和运行时。模型从训练到部署，就像在不同国家之间穿梭，得不断换护照、办签证。我当年在项目里就吃过这个亏——算法同学给了一个PyTorch的模型，我们NPU的编译器只认TensorFlow的PB文件，中间折腾了整整两周，光格式转换就写了一堆屎山代码。
2017年，微软和Facebook牵头搞了ONNX（Open Neural Network Exchange）。说白了，就是给AI模型定一个“通用语言”。你拿PyTorch训练完，导出成ONNX；我拿TensorFlow训练完，也导出成ONNX。到了我们NPU这边，只要解析ONNX就行了，不用管你背后是啥框架。这就像全世界都讲普通话，交流成本瞬间降下来了。
核心价值：ONNX不是一个新的训练框架，它只是一个中间表示（Intermediate Representation, IR）。它的使命是“一次导出，到处运行”。
4.2 ONNX算子集定义：你必须要懂的OpSet
ONNX定义了一套标准的算子集合，叫做OpSet。每个版本都会新增或废弃一些算子。比如ONNX opset 11支持了ScatterND，opset 13支持了Softmax的改进版本。我建议你在导出模型时，尽量使用较新的opset版本，但别太激进——太新的版本，有些NPU编译器可能还没适配。
常见的ONNX算子包括：
- 基础数学运算：Add、Sub、Mul、Div、MatMul
- 激活函数：Relu、Sigmoid、Tanh、Softmax
- 张量操作：Reshape、Transpose、Concat、Split、Slice
- 卷积/池化：Conv、ConvTranspose、MaxPool、AveragePool
- 归一化：BatchNormalization、LayerNormalization
- 控制流：If、Loop、Scan
嗯，这里要注意：ONNX的算子粒度比PyTorch要粗一些。比如PyTorch里的F.relu，在ONNX里就是Relu节点。但有些PyTorch的自定义操作，ONNX原生不支持，这时候就需要你手动写一个自定义算子，或者用torch.onnx.export的custom_opsets参数去注册。
避坑指南：我曾经在导出Transformer模型时，发现PyTorch的nn.MultiheadAttention在ONNX里被拆成了十几个小算子，导致NPU推理性能下降30%。后来我手动重写了注意力机制，用torch.onnx.export的operator_export_type参数强制合并算子，才把性能拉回来。所以，导出后一定要用onnxruntime或者onnxsim做一次模型优化。
4.3 模型结构解析：读懂ONNX的“骨架”
一个ONNX模型，本质上是一个计算图。它由三部分组成：
- GraphProto：整个计算图的描述，包含所有节点、输入、输出、初始值。
- NodeProto：每个算子节点，包含算子类型、输入输出张量名、属性参数。
- TensorProto：权重和偏置等常量数据，以二进制形式存储。
你可以用onnx.load()加载模型，然后用onnx.helper.printable_graph()打印出可读的图结构。我个人习惯用netron这个可视化工具，它能把ONNX模型渲染成一张漂亮的流程图，一眼就能看出网络结构。
下面我用SVG画一张ONNX模型结构的示意图，帮你理解它的层次关系：
你看，整个结构是树形的。最顶层是ModelProto，它包含一个GraphProto，而GraphProto里又包含多个NodeProto（算子节点）、TensorProto（权重数据）和ValueInfoProto（输入输出信息）。每个NodeProto就是一个具体的算子，比如Conv、Relu、Pool等。
4.4 如何将PyTorch/TF模型导出为ONNX
这部分是实操重点。我分别讲PyTorch和TensorFlow的导出方法。
4.4.1 PyTorch模型导出
PyTorch官方提供了torch.onnx.export函数。基本用法如下：
import torch
import torch.onnx
# 假设你有一个训练好的模型
model = MyModel()
model.eval()
# 创建一个虚拟输入，形状要和实际输入一致
dummy_input = torch.randn(1, 3, 224, 224)
# 导出ONNX
torch.onnx.export(
model, # 模型
dummy_input, # 虚拟输入
"model.onnx", # 输出文件名
export_params=True, # 是否导出权重
opset_version=13, # ONNX算子集版本
do_constant_folding=True, # 是否做常量折叠优化
input_names=['input'], # 输入张量名
output_names=['output'], # 输出张量名
dynamic_axes={
'input': {0: 'batch_size'}, # 动态batch
'output': {0: 'batch_size'}
}
)
这里有几个关键参数：
- opset_version：我建议用13或以上，支持更多算子。但别用太新，比如17，有些NPU工具链还没跟上。
- dynamic_axes：如果你的模型需要支持可变batch size，一定要设置这个。否则导出的模型只能跑固定batch。
- do_constant_folding：建议开启，能减少一些冗余计算节点。
注意：PyTorch导出ONNX时，如果模型里有if条件分支或者for循环，可能会导出失败。我曾经遇到一个模型用了torch.where，导出后ONNX图里多了一个If节点，NPU编译器直接报错。解决办法是用torch.jit.script先把模型转成TorchScript，再导出ONNX。
4.4.2 TensorFlow模型导出
TensorFlow这边，推荐用tf2onnx工具。先安装：
pip install tf2onnx
然后导出：
import tensorflow as tf
import tf2onnx
# 加载你的TF模型
model = tf.keras.models.load_model('my_model.h5')
# 指定输入签名
spec = (tf.TensorSpec((None, 224, 224, 3), tf.float32, name="input"),)
# 导出ONNX
output_path = "model.onnx"
model_proto, _ = tf2onnx.convert.from_keras(
model,
input_signature=spec,
opset=13,
output_path=output_path
)
对于TF SavedModel格式，可以用命令行：
python -m tf2onnx.convert \
--saved-model ./saved_model \
--output model.onnx \
--opset 13
我个人经验是，TF模型导出ONNX时，动态轴的处理比PyTorch麻烦。TF默认是静态图，如果你需要动态batch，必须在TensorSpec里把None传进去。另外，TF的一些高级API（比如tf.keras.layers.MultiHeadAttention）导出时可能会被拆成很多小算子，建议导出后用onnxsim做一次简化。
验证导出结果：导出后，一定要用onnxruntime跑一次推理，对比输出是否和原框架一致。我习惯写一个脚本，分别用PyTorch/TF和ONNX Runtime跑同一个输入，然后比较输出的余弦相似度。如果相似度低于0.999，说明导出有问题，需要排查。
好了，关于ONNX的诞生背景、算子集、模型结构以及导出方法，就讲到这里。记住一句话：ONNX是AI芯片部署的“通用语言”，掌握它，你就能在框架和芯片之间自由穿梭。