4、ONNX算子与数据类型：ONNX标准算子集、自定义算子、数据类型映射、算子版本兼容性

好，咱们今天聊聊ONNX里最核心的两个东西——算子和数据类型。

说实话，我刚开始接触ONNX的时候，觉得这玩意儿就是个模型交换格式嘛，把模型转过去就完事了。结果一跑起来，各种报错，什么“Unsupported Operator”、“Type Mismatch”，搞得我头大。后来我才明白，不懂算子和数据类型，你连模型都转不对，更别提优化了。

4.1 ONNX标准算子集：你的工具箱

ONNX定义了一套标准算子，说白了就是给AI模型里的各种操作起了个统一的名字。比如卷积叫Conv，全连接叫Gemm，ReLU就叫Relu。这套算子集是跨框架的，不管你是PyTorch、TensorFlow还是飞桨，转成ONNX后，大家都用同一套名字。

我个人习惯把算子分成几类：

  • 张量操作：Reshape、Transpose、Concat、Split。这些是“搬砖”的，不改变数值，只改变形状或排列。
  • 数学运算：Add、Mul、MatMul、Sqrt。加减乘除，矩阵乘法，这些是“算数”的。
  • 激活函数：Relu、Sigmoid、Tanh。给模型加非线性的。
  • 归一化：BatchNormalization、InstanceNormalization。稳定训练用的。
  • 池化与卷积：Conv、MaxPool、AveragePool。图像处理的主力。
  • 控制流：If、Loop、Scan。这个比较高级，动态图里常用。

嗯，这里要注意：ONNX的算子集是不断扩充的。比如早期的版本没有Resize算子，后来为了支持上采样才加进去。所以你在转模型时，最好确认一下目标ONNX版本支持哪些算子。

核心观点：标准算子集是ONNX的“通用语言”。你写的模型最终都会被拆解成这些基本操作。理解它们，你才能看懂ONNX模型在干什么。

4.2 自定义算子：当标准工具箱不够用时

我在项目中遇到过一种情况：模型里用了一个非常小众的操作，比如某种自定义的注意力机制，ONNX标准算子集里根本没有。怎么办？硬转肯定报错。

这时候就需要自定义算子（Custom Op）了。说白了，就是你自己写一个算子的实现，注册到ONNX Runtime里，让它认识这个操作。

自定义算子的流程大致是这样的：

  1. 定义算子：在ONNX模型中，用一个自定义的节点（比如“MyCustomOp”）代替原来的操作。
  2. 注册实现：在ONNX Runtime里，用C++或Python写一个这个算子的前向计算逻辑，然后注册进去。
  3. 加载运行：运行时，ONNX Runtime遇到这个自定义节点，就会调用你写的实现。

举个例子，假设你有一个自定义的“软阈值”操作：

# 伪代码，展示自定义算子的注册思路
import onnxruntime as ort

# 定义一个自定义算子的实现
class SoftThresholdKernel:
    def compute(self, input_tensor, threshold):
        return np.sign(input_tensor) * np.maximum(np.abs(input_tensor) - threshold, 0)

# 注册到ONNX Runtime（简化示意）
ort.register_custom_op("MyDomain::SoftThreshold", SoftThresholdKernel)

警告：自定义算子虽然灵活，但会破坏ONNX的跨平台性。你想想看，你的模型在A机器上能跑，换到B机器上，如果B机器没注册这个算子，直接就崩了。所以，能不用就不用，实在没办法再用。

我曾经为了一个自定义的池化操作，折腾了整整两天。最后发现其实可以用标准算子组合实现，性能还更好。嗯，有时候“偷懒”用标准算子反而是最优解。

4.3 数据类型映射：别让精度坑了你

数据类型映射，听起来很学术，其实就一句话：不同框架里的数据类型，在ONNX里怎么对应？

比如PyTorch里是torch.float32，TensorFlow里是tf.float32，到了ONNX里统一叫FLOAT。再比如PyTorch的torch.int64，ONNX里叫INT64。

我整理了一个常用的映射表，你直接拿去用：

PyTorch TensorFlow ONNX 描述
torch.float32 tf.float32 FLOAT 单精度浮点，最常用
torch.float64 tf.float64 DOUBLE 双精度，精度高但慢
torch.int32 tf.int32 INT32 32位整数
torch.int64 tf.int64 INT64 64位整数，索引常用
torch.bool tf.bool BOOL 布尔类型
torch.uint8 tf.uint8 UINT8 无符号8位，图像常用

为什么会搞混？我遇到过最坑的一次，是PyTorch模型里用了torch.int64做索引，转成ONNX后，在某个推理引擎上只支持INT32。结果模型跑起来，索引直接溢出，结果全错了。从那以后，我转模型前都会先检查一遍数据类型。

小技巧：如果你不确定数据类型，可以在导出ONNX时显式指定。比如在PyTorch里用torch.onnx.export时，可以传入一个dummy input，确保它的dtype是你想要的。

4.4 算子版本兼容性：版本号里藏着坑

算子版本兼容性，说白了就是：同一个算子，不同ONNX版本，行为可能不一样。

举个例子，ONNX opset 10里的Resize算子和opset 11里的Resize算子，参数定义就变了。opset 10用的是“scale”参数，opset 11改成了“sizes”参数。如果你用opset 10的模型在opset 11的Runtime上跑，Runtime会尝试做兼容，但有时候会失败。

我个人习惯的做法是：

  • 统一版本：训练和推理都用同一个ONNX opset版本。比如训练时用opset 15，推理时Runtime也支持opset 15。
  • 向下兼容：如果必须用低版本模型，确保Runtime支持该版本。ONNX Runtime一般会向前兼容，但不敢保证100%。
  • 查看文档：每个算子的版本变更，ONNX官方文档都有详细说明。我每次升级opset版本，都会去翻一下变更日志。

避坑指南：我曾经把一个opset 11的模型部署到只支持opset 10的Runtime上，结果Resize算子直接报错。最后我不得不重新导出模型，把opset降到10。所以，导出前先确认目标Runtime的opset支持范围。

嗯，总结一下。算子是你的工具，数据类型是你的材料，版本是你的施工规范。这三样搞明白了，ONNX模型部署就成功了一半。下一章咱们聊聊模型转换的实战技巧，到时候我会拿几个真实案例来讲。

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