ONNX 算子与图结构：版本、自定义与拓扑

这一章咱们要啃的，是 ONNX 模型最核心的骨架——算子和图结构。说白了，模型就是个有向无环图，节点是算子，边是张量。你只有搞懂了算子集怎么管理版本、怎么塞进自定义算子，以及图结构怎么遍历、怎么切片，才能真正掌控模型的优化流程。

算子集版本：版本号背后的学问

ONNX 算子集有个版本号，比如 opset 11、opset 15。很多人觉得它就是个数字，其实不是。它直接决定了你的模型能不能被推理引擎正确加载。

我记得有一次项目组从 PyTorch 导出模型，用的 opset 17，结果对方那个 TensorRT 5 的环境最高只支持 opset 11。模型死活加载不了。后来只能回退到 opset 11 重新导出，才解决问题。

查看算子集版本很简单，看模型元数据就行：

import onnx

model = onnx.load("model.onnx")
print(model.opset_import[0].version)  # 输出算子集版本号

版本演进背后是算子的新增和修改。比如 opset 11 加入了 ScatterND，opset 13 里 Softmax 的 axis 参数变成必选。这些细节不注意，转换时就会报错。

算子集版本演进对照表

算子集版本	关键变化	推荐场景
opset 7	基础算子定型，支持 RNN/LSTM	老旧模型兼容
opset 9	加入 ConstantOfShape、Upsample	图像类模型常见
opset 11	加入 ScatterND、ScatterElements；Resize 改参数	通用兼容性最优
opset 13	Softmax 修改；加入 GridSample	Transformer、分割模型
opset 15	Resize 参数调整，加入 BitShift	前沿模型导出

自定义算子：补齐 ONNX 的短板

标准算子库覆盖了 90% 的常见操作，但总有那么几个自定义 op 你绕不过去——比如你写了个特殊的激活函数，或者用到了一个第三方库里的高效算子。

这时候就需要自定义算子。说白了就是先在 ONNX 里注册一个你的新算子，然后在推理引擎里实现它的计算逻辑。

在 ONNX 里注册自定义算子，需要两步：定义符号函数，注册适配器。

import torch
import torch.onnx.symbolic_helper as sym_help
from torch.onnx import register_custom_op_symbolic

# 第一步：定义符号函数
def my_custom_relu(g, input, alpha):
    # 映射到 ONNX 的自定义算子
    return g.op("my::CustomRelu", input, alpha_i=alpha)

# 第二步：注册适配器
register_custom_op_symbolic("::my_custom_relu", my_custom_relu, 1)

# 在模型中使用
class MyModel(torch.nn.Module):
    def forward(self, x):
        return torch.ops.my_custom_relu(x, alpha=0.1)

在 TensorRT 里加载自定义算子，就得写插件（Plugin）。我曾经写过一个 FFT 算子插件，调了整整两天才跑通。关键点是 input 和 output 的维度推导一定要测全，不然推理时会崩在奇怪的地方。

图遍历与拓扑排序：搞懂计算顺序

ONNX 模型是一个有向无环图（DAG）。图中每个节点代表一个算子，边代表数据流向。你要做优化，第一步就是遍历并理解这个图的结构。

图遍历分两种：广度优先（BFS）和深度优先（DFS）。我记得做模型剪枝时，就是用 BFS 从输出节点往回找，逐步剥离冗余的层。

拓扑排序是 DAG 优化的基础。它保证了：如果算子 A 的输入依赖算子 B 的输出，那么 B 一定在 A 之前执行。这是所有推理引擎的底层基础。

import onnx
from collections import deque

model = onnx.load("model.onnx")
graph = model.graph

# 拓扑排序：Kahn 算法
in_degree = {node.name: 0 for node in graph.node}
# 建立邻接表和入度表
# ...（省略构建过程，假设已经构建好）

queue = deque([n for n, d in in_degree.items() if d == 0])
sorted_nodes = []

while queue:
    node = queue.popleft()
    sorted_nodes.append(node)
    for neighbor in adj_list[node]:
        in_degree[neighbor] -= 1
        if in_degree[neighbor] == 0:
            queue.append(neighbor)

# sorted_nodes 就是拓扑排序后的算子列表

子图与函数：封装与复用

子图，就是计算图的一部分。你用它对模型做剪枝、量化或者融合优化。TensorRT 的 layer fusion 本质上就是在子图上做模式匹配，把多个小算子合并成一个更大的算子。

函数（Function）是 ONNX 1.6 之后引入的新特性。它让你定义可复用的算子组合。比如你有个“conv + bn + relu”的组合，就可以定义成一个函数，在多处复用。

我做过一个项目，把 68 个重复的卷积激活块定义成函数，模型体积直接减少了 30%，导出速度也快了。

# 定义一个函数：ConvBNReLU
func_proto = onnx.helper.make_function(
    domain="my.domain",
    fname="ConvBNReLU",
    inputs=["X", "W", "B", "gamma", "beta", "mean", "var"],
    outputs=["Y"],
    nodes=[
        # Conv 节点
        onnx.helper.make_node("Conv", ["X", "W", "B"], ["conv_out"], 
                              kernel_shape=[3,3], pads=[1,1,1,1]),
        # BN 节点
        onnx.helper.make_node("BatchNormalization", 
                              ["conv_out", "gamma", "beta", "mean", "var"],  
                              ["bn_out"], epsilon=1e-5),
        # ReLU 节点
        onnx.helper.make_node("Relu", ["bn_out"], ["Y"])
    ],
    opset_imports=[onnx.helper.make_opsetid("", 13)]
)

# 在图中使用函数
graph.node.append(
    onnx.helper.make_node("my.domain.ConvBNReLU", 
                          ["X", "W", "B", "gamma", "beta", "mean", "var"],
                          ["Y"], domain="my.domain")
)

一张图概括本章知识体系

避坑指南与个人经验

关于算子版本，我常用的做法是定义一个字典，把目标推理引擎支持的 opset 范围记下来。导出模型前先查这个字典，能省不少排查时间。

自定义算子最容易出坑的地方是 shape inference。如果你不实现它，TensorRT 或者 ONNXRuntime 会认为输出形状未知，然后就可能报错或者走很慢的 fallback 路径。

图拓扑排序在我做算子融合时帮了大忙。那时候需要把相邻的 Conv 和 BN 合并，用拓扑序就能从左到右一次扫描，不用来回看图。