4、自定义算子支持：通过 torch.onnx.register_custom_op_symbolic 注册新算子

说实话，模型转换中最让人头疼的，就是遇到 PyTorch 原生不支持的算子。

我刚开始做 ONNX 部署时，就栽过这个跟头。一个看起来很简单的自定义操作，导出时直接报错，卡了我整整两天。后来才发现，原来 PyTorch 早就给我们留了后门——torch.onnx.register_custom_op_symbolic。

说白了，这个函数就是让你告诉 ONNX：「嘿，我这个自定义操作，应该映射成 ONNX 里的哪个算子，参数怎么传。」

4.1 什么时候需要自定义算子？

先说说我遇到过的几种典型场景：

  • 自己写的自定义函数：比如一个特殊的激活函数，PyTorch 里用 torch.autograd.Function 实现的
  • 第三方库里的操作：像某些检测模型里的 NMS（非极大值抑制），PyTorch 原生导出不支持
  • 性能优化后的融合算子：比如把 Conv+BN+ReLU 合并成一个算子，推理框架不认识

嗯，这里要注意：不是所有自定义操作都需要注册。如果你的操作可以用 PyTorch 原生算子组合实现，那 ONNX 导出时通常会自动展开。只有那些「黑盒」操作，才需要你手动告诉 ONNX 怎么处理。

4.2 注册的基本流程

我个人习惯把注册过程分成三步：

  1. 定义符号函数（symbolic）：描述这个操作在 ONNX 里长什么样
  2. 调用 register_custom_op_symbolic：把 PyTorch 算子和符号函数绑定
  3. 正常导出：PyTorch 遇到这个算子时，会自动调用你注册的符号函数

来看一个最简单的例子。假设我写了一个自定义的「平方加一」操作：

import torch
import torch.onnx

# 1. 定义符号函数
def my_square_plus_one_symbolic(g, input):
    """把自定义操作映射成 ONNX 的 Pow + Add"""
    # g 是 ONNX 的图构建器
    # 先计算平方
    squared = g.op("Pow", input, g.op("Constant", value_t=torch.tensor(2.0)))
    # 再加 1
    result = g.op("Add", squared, g.op("Constant", value_t=torch.tensor(1.0)))
    return result

# 2. 注册算子
torch.onnx.register_custom_op_symbolic(
    symbolic_name="mylib::square_plus_one",  # 命名空间::操作名
    symbolic_fn=my_square_plus_one_symbolic,
    opset_version=11  # 指定 ONNX opset 版本
)

# 3. 使用
class MyModel(torch.nn.Module):
    def forward(self, x):
        return torch.ops.mylib.square_plus_one(x)

model = MyModel()
dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(model, dummy_input, "model.onnx", opset_version=11)

关键点：符号函数的第一个参数永远是 g（ONNX 图构建器），后面的参数要和 PyTorch 算子的输入一一对应。

4.3 符号函数怎么写？

写符号函数，说白了就是在用 Python 拼 ONNX 的图结构。我常用的几个技巧：

  • 常量用 g.op("Constant", value_t=...)：ONNX 里常量需要显式创建
  • 属性用 keyword 参数传：比如 g.op("Conv", input, weight, kernel_shape=[3,3])
  • 多输出用 g.op 返回 tuple：ONNX 支持一个节点多个输出

举个例子，我曾经需要导出一个自定义的 ROI Pooling 操作：

def my_roi_pool_symbolic(g, features, rois, output_size, spatial_scale):
    """自定义 ROI Pooling 映射到 ONNX 的 RoiAlign"""
    return g.op(
        "RoiAlign",
        features,
        rois,
        output_height_i=output_size[0],  # 属性用 _i 后缀表示 int
        output_width_i=output_size[1],
        spatial_scale_f=spatial_scale,   # 属性用 _f 后缀表示 float
        sampling_ratio_i=2,
        coordinate_transformation_mode_s="half_pixel"  # _s 表示 string
    )

小技巧：ONNX 的属性后缀规则：_i 整数，_f 浮点数，_s 字符串，_t 张量，_is 整数列表。记不住没关系，我每次写的时候都要翻文档确认一下。

4.4 处理多输入和多输出

你想想看，实际项目中哪有那么简单的一进一出？我遇到过最复杂的一个自定义算子，有 5 个输入、3 个输出。

多输入好办，直接在符号函数参数里列出来就行。多输出呢？ONNX 里一个节点可以返回多个值，用 tuple 接收：

def my_detection_symbolic(g, scores, boxes, nms_thresh, score_thresh):
    """自定义检测后处理，返回检测框、分数、类别"""
    # 假设内部调用了 ONNX 的 NonMaxSuppression
    selected_indices = g.op(
        "NonMaxSuppression",
        boxes, scores,
        max_output_boxes_per_class_i=100,
        iou_threshold_f=nms_thresh,
        score_threshold_f=score_thresh
    )
    # 用 Gather 取出结果
    # ... 省略中间处理逻辑
    return det_boxes, det_scores, det_classes  # 返回多个输出

注意：ONNX 导出时，如果你的自定义算子返回多个值，调用方也要用多个变量接收。我曾经因为少写了一个返回值，调试了整整一个下午。

4.5 避坑指南

嗯，这部分是我最想分享的。踩过的坑多了，自然就记住了。

  • 命名空间不能乱用：torch.onnx.register_custom_op_symbolic 的第一个参数格式是 命名空间::操作名。命名空间建议用项目名或库名，别用 torch、onnx 这些保留字。
  • opset 版本要匹配：你注册时指定的 opset 版本，必须和导出时的 opset_version 一致。我习惯统一用 11 或 13，这两个版本支持比较稳定。
  • 动态 shape 要小心：如果你的算子输出 shape 依赖于输入数据（比如 NMS 的输出框数量不确定），ONNX 导出时可能会报 shape 推断错误。这时候可以试试 torch.onnx.export(..., dynamic_axes=...)。
  • 先验证再部署：注册完一定要用 onnxruntime 跑一遍推理，看看结果对不对。我曾经注册了一个算子，导出没报错，但推理结果全是 NaN。

4.6 实战：注册一个自定义激活函数

最后，我带你走一遍完整的流程。假设我们有一个自定义的激活函数 my_swish：

import torch
import torch.onnx

# 定义 PyTorch 算子
class MySwishFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        ctx.save_for_backward(x)
        return x * torch.sigmoid(x)
    
    @staticmethod
    def backward(ctx, grad_output):
        x, = ctx.saved_tensors
        sigmoid_x = torch.sigmoid(x)
        return grad_output * (sigmoid_x + x * sigmoid_x * (1 - sigmoid_x))

# 注册符号函数
def my_swish_symbolic(g, x):
    """Swish = x * sigmoid(x)"""
    sigmoid = g.op("Sigmoid", x)
    return g.op("Mul", x, sigmoid)

torch.onnx.register_custom_op_symbolic(
    "mylib::my_swish", my_swish_symbolic, opset_version=11
)

# 测试导出
class SwishModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = torch.nn.Linear(10, 5)
    
    def forward(self, x):
        x = self.fc(x)
        return MySwishFunction.apply(x)  # 使用自定义算子

model = SwishModel()
dummy = torch.randn(1, 10)
torch.onnx.export(model, dummy, "swish_model.onnx", opset_version=11)
print("导出成功！")

你看，其实没那么复杂。核心就是三步：写符号函数、注册、导出。一旦你掌握了这个套路，再奇怪的算子也能搞定。

我个人建议，刚开始接触时先拿简单的数学操作练手，比如自定义的加法、乘法。等熟悉了 g.op 的用法，再去挑战复杂的检测、分割算子。一步一步来，别着急。

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