4、ONNX算子详解：常见算子（Conv、Relu、Gemm、BatchNormalization等）、算子版本兼容性、自定义算子
好，咱们进入第四章。这一章我打算跟你好好聊聊ONNX里的算子。
算子是什么？说白了，就是神经网络里的基本操作单元。你搭建的每一层网络，最后都会映射成一个个ONNX算子。我刚开始接触ONNX时，觉得这玩意儿就是个中间格式，没啥好学的。后来踩了坑才发现，不懂算子，你连模型为啥转失败都搞不清楚。
4.1 常见算子详解
咱们先看几个最常用的。这几个算子，你几乎在每个模型里都能见到。
4.1.1 Conv（卷积）
Conv算子，做卷积运算的。它有几个关键属性：dilations（膨胀系数）、group（分组数）、kernel_shape（卷积核尺寸）、pads（填充）、strides（步长）。
嗯，这里要注意一点。ONNX里的Conv，输入数据的格式是 [N, C, H, W]。N是batch size，C是通道数，H和W是高度和宽度。权重格式是 [M, C/group, kH, kW]。M是输出通道数。
我在项目中遇到过一个问题：用PyTorch训练的模型，转ONNX后推理结果不对。查了半天，发现是PyTorch默认的padding模式和ONNX的不完全一致。PyTorch的padding='same'在某些配置下，算出来的填充值跟ONNX的auto_pad='SAME_UPPER'有细微差别。所以，我建议你转模型时，尽量显式指定padding的具体数值，别偷懒用字符串。
4.1.2 Relu（线性整流单元）
Relu算子，简单到不能再简单：y = max(0, x)。它只有一个输入和一个输出，没有额外属性。
但你别小看它。我曾经在部署一个MobileNet模型时，发现推理速度比预期慢。排查后发现，某些推理框架对Relu做了融合优化——把Relu合并到前一个Conv或Gemm里。如果模型里Relu是单独一个节点，就没法触发这个优化。所以，我建议你在训练时，尽量用Conv+Relu这种组合，而不是分开写。
4.1.3 Gemm（通用矩阵乘法）
Gemm算子，就是全连接层。它的公式是：Y = alpha * A * B + beta * C。其中alpha和beta是标量，默认都是1.0。transA和transB控制是否对矩阵进行转置。
你想想看，为什么要有transA和transB？因为不同框架存储权重的方式不一样。PyTorch的Linear层，权重形状是[out_features, in_features]，而ONNX的Gemm默认transB=1，也就是把权重转置后再乘。如果你从TensorFlow转过来，可能就得调整这个标志位。
我个人习惯是，转完模型后，用onnxruntime跑一下推理，跟原始框架的结果对比。差一个1e-5以内算正常，差太多就要检查Gemm的转置标志了。
4.1.4 BatchNormalization（批归一化）
BatchNormalization，简称BN。它在推理时，公式是：y = (x - mean) / sqrt(var + epsilon) * scale + bias。
ONNX里的BN算子，有5个输入：X、scale、bias、mean、var。还有两个属性：epsilon（防止除零的小数）和momentum（动量，推理时其实用不到）。
这里有个坑。不同版本的ONNX，BN算子的行为有变化。在opset 9之前，BN算子的输出有4个（除了归一化后的数据，还有均值、方差等的更新值）。从opset 9开始，输出简化成了1个。如果你用旧版工具链转模型，可能会遇到输出数量不匹配的问题。
核心要点：转模型时，明确指定opset版本。我一般用opset 11或13，兼容性比较好。
4.2 算子版本兼容性
算子版本兼容性，说白了就是：你的ONNX模型，能不能在目标推理引擎上跑起来。
ONNX每个版本，都会新增或修改一些算子。比如，Resize算子在opset 10和opset 11中，坐标变换模式就不一样。opset 10用的是half_pixel，opset 11改成了asymmetric。如果你用opset 11的Resize，但推理引擎只支持opset 10，那就会报错。
我建议你遵循一个原则：用低版本，不用高版本。除非你确定目标平台支持。比如，你要部署到手机端，那最好用opset 10或11。云端部署的话，opset 13、15甚至17都可以试试。
怎么查看模型用了哪些算子？用onnx库就行：
import onnx
model = onnx.load('model.onnx')
opsets = model.opset_import
for op in opsets:
print(f"Domain: {op.domain}, Version: {op.version}")
输出结果里，domain为""或"ai.onnx"的是标准算子，domain为"ai.onnx.ml"的是机器学习算子，还有自定义的domain。
小技巧：如果你不确定某个算子是否被支持，可以用onnxruntime的get_available_providers()和get_provider_options()来查看。或者直接写个测试脚本，跑一下包含该算子的简单模型。
4.3 自定义算子
有时候，标准算子不够用。比如，你实现了一个特殊的激活函数，或者一个高效的注意力机制。这时候，就需要自定义算子。
自定义算子分两步：
- 注册算子：在ONNX中定义算子的名称、输入输出、属性。
- 实现算子：在推理引擎中实现算子的计算逻辑。
举个例子。假设我写了一个MyActivation算子，公式是y = x^2 + 1。在Python中注册它：
import onnx
from onnx import helper, TensorProto
# 定义算子
my_op = helper.make_node(
'MyActivation', # 算子名称
inputs=['X'], # 输入
outputs=['Y'], # 输出
domain='my.custom', # 域，避免跟标准算子冲突
)
# 创建图
graph = helper.make_graph(
[my_op],
'my_graph',
[helper.make_tensor_value_info('X', TensorProto.FLOAT, [1, 3, 224, 224])],
[helper.make_tensor_value_info('Y', TensorProto.FLOAT, [1, 3, 224, 224])],
)
model = helper.make_model(graph, opset_imports=[helper.make_opsetid('', 13), helper.make_opsetid('my.custom', 1)])
然后，在推理引擎里实现它。以ONNX Runtime为例，你需要写一个C++或Python的算子内核，并注册到OrtCustomOpDomain里。
警告：自定义算子会破坏模型的跨平台兼容性。你想想看，如果模型里用了my.custom.MyActivation，那换一个推理引擎，就得重新实现一遍。所以，我建议你：能用标准算子组合实现的，就别自定义。实在不行，再考虑自定义。
我曾经在一个项目里，为了优化一个特殊的分组卷积，写了个自定义算子。结果部署到不同硬件时，每个平台都要写一遍实现，维护成本极高。后来我改用标准Conv+Split+Concat组合，虽然速度慢了一点，但省去了大量适配工作。
嗯，这一章的内容就这些。算子这块，你只要把常见的几个搞明白，再掌握版本兼容性的判断方法，基本就能应对大部分场景了。自定义算子，知道怎么用就行，别轻易上生产环境。