目录前言为什么要看懂算子注册一、ONNX 导出本质上是翻译二、4 种算子不支持的工况工况 1词典缺条目PyTorch 有、ONNX 标准也有、只是导出器没映射工况 2目标语言没这个词torchvision 有、ONNX 标准没有目标语言没这个词torchvision 有、ONNX 标准没有工况 3自研算子autograd.Function工况 4opset 版本问题用了旧版词典三、4 种注册写法对应不同场景四、关键辨析新手必踩的 3 个坑4.1 两个算子名不是一回事4.2 导出阶段 vs 构建阶段上半场与下半场4.3 parse_args 的类型字母表五、参数表的来源模仿官方格式5.1 底层 C 签名的具体实现5.2 陌生算子的签名查询5.2.1 **torchvision算子通过算子名称直接拿底层签名**5.2.2 怎么知道算子在哪个命名空间列出所有算子5.3 注册算子时模仿的官方格式映射函数六、 ONNX 算子文档获取与理解Operators.md 全攻略方式 1Gitee 镜像最直接内容与 GitHub 完全一致方式 2ONNX 官网带搜索体验更好方式 3本地源码仓库如果有 onnx 源码七Netron 可视化看已有模型用八、遇到导出报错的标准排查流程实战总结九、总结前言为什么要看懂算子注册做模型部署的人几乎每天都在跟这条流水线打交道而**算子不支持是这条流水线上最高频的拦路虎**。你在 PyTorch 里写得好好的一个模型一执行torch.onnx.export就报一堆看不懂的错好不容易导出成功了喂给trtexec构建 engine 又报错。本文以及后续几篇文章要解决的核心问题就一个当算子不认识时到底该怎么办本文将给出 4 种工况针对不兼容的4种工况给出相应的解决方案最后附上去哪里查算子文档的保姆级清单。一、ONNX 导出本质上是翻译PyTorch 导出的过程本质上是一个翻译过程PyTorch 的计算图是用ATen 算子a Tensor LibraryPyTorch 底层的 C 张量运算库写成的——可以理解为中文ONNX 的计算图是用ONNX 标准算子写成的——可以理解为英文导出器torch.onnx.export就是一个翻译官它手里有一本词典记录了每个 ATen 算子中文词对应哪个 ONNX 算子英文词。这本词典就是 PyTorch 源码里的torch/onnx/symbolic_opset*.py文件。每一个词条叫符号函数symbolic function。如果词典里没有某个中文词的条目翻译官就卡壳了于是导出报错。注册算子就是手动给这本词典补一条新条目“这个中文词aten::xxx翻译成英文词Xxx这样翻。”二、4 种算子不支持的工况把报错归归类你会发现算子不支持其实分 4 种完全不同的情况处理方式天差地别。工况 1词典缺条目PyTorch 有、ONNX 标准也有、只是导出器没映射代表torch.asinh示例工程 3.4 的sample_asinh.py现象直接导出报错。根因PyTorch 有torch.asinh()ATen 里有aten::asinh✅ONNX 标准从opset 9起就有Asinh算子 ✅但是用的 PyTorch 版本的symbolic_opset*.py里没写aten::asinh→Asinh的映射条目 ❌就像中文有asinh这个词、英文也有Asinh这个词但翻译官的词典漏了这一条。解决fromtorch.onnximportregister_custom_op_symbolicdefasinh_symbolic(g,input,*,outNone):returng.op(Asinh,input)# 左边是 ONNX 算子名register_custom_op_symbolic(aten::asinh,asinh_symbolic,12)# ↑ PyTorch 算子名就 3 行核心代码手动补一条翻译条目。导出时遇到aten::asinh就知道要生成 ONNX 的Asinh节点。工况 2目标语言没这个词torchvision 有、ONNX 标准没有目标语言没这个词torchvision 有、ONNX 标准没有代表torchvision.ops.DeformConv2d现象直接导出报错。根因torchvision 有DeformConv2d底层有完整 CUDA 实现 ✅ONNX 标准里从来没有可变形卷积这个算子❌就像中文有内卷这个词但英文里根本没有对应的标准词汇——你不能把它翻译成一个不存在的词。解决sample_deformable_conv_register.pyfromtorch.onnximportregister_custom_op_symbolicfromtorch.onnx.symbolic_helperimportparse_argsparse_args(v,v,v,v,v,i,i,i,i,i,i,i,i,none)defdcn_symbolic(g,input,weight,offset,mask,bias,stride_h,stride_w,pad_h,pad_w,dil_h,dil_w,n_weight_grps,n_offset_grps,use_mask):returng.op(custom::deform_conv2d,input,offset)# 注意自定义 domainregister_custom_op_symbolic(torchvision::deform_conv2d,dcn_symbolic,12)注意这里的关键区别g.op(custom::deform_conv2d, ...)用了自定义 domaincustom::而不是g.op(DeformConv2d, ...)——因为 ONNX 标准里没这个算子你不能凭空编一个标准算子。这意味着什么导出成custom::deform_conv2d后问题并没有真正解决完只是上半场onnxruntime 默认跑不了不认识这个算子TensorRT 也跑不了必须自己写一个TensorRT PluginC/CUDA kernel来实现可变形卷积的前向计算也就是说我们虽然可以顺利导出onnx文件但实质上这部分是个“空壳”这就是下半场后面有时间再更新相关的文章敬请期待。提醒可变形卷积DCN在目标检测模型CenterPoint、FCOS3D 等中非常常见部署时几乎都要写 TensorRT Plugin。这是部署工程师的核心技能之一。工况 3自研算子autograd.Function代表自己写的一个torch.autograd.Function算子。现象这一组不是报错而是导出结果不符合预期——你写的自定义算子被**拆开inline**成了一串基础节点。classCustomOp(torch.autograd.Function):staticmethoddefforward(ctx,x:torch.Tensor)-torch.Tensor:ctx.save_for_backward(x)xx.clamp(min0)returnx/(1torch.exp(-x))导出后的onnx文件网络结构是这样的根因如果用torch.autograd.Function写了一个自定义算子导出器在不注册 symbolic时会追踪 forward 里的每一个基础算子最终导出的 ONNX 图是Clip → Exp → Add → Div一串节点。图能跑但不是一个整体算子。解决在 Function 子类里定义symbolic静态方法importtorchclassCustomOp(torch.autograd.Function):staticmethoddefsymbolic(g,x):returng.op(custom_domain::customOp2,x)# 导出成单个自定义节点staticmethoddefforward(ctx,x):xx.clamp(min0)returnx/(1torch.exp(-x))customOpCustomOp.apply导出后的onnx文件网络结构是这样的什么时候用哪种如果自定义算子的计算能用 ONNX 标准算子组合表达、且性能没问题 → 不注册 symbolic让它 inline 展开就行如果需要在推理引擎里对算子做特殊优化融合、写 Plugin→ 必须注册 symbolic 导出为单个节点否则引擎看不到这是一个整体没法优化。工况 4opset 版本问题用了旧版词典现象同一个算子换个 opset 版本就能导出或换个版本就报错。根因ONNX 标准在演进每个 opset 版本会新增 / 废弃 / 修改算子。常见例子算子引入版本说明Asinhopset 9更低版本就没有LayerNormalizationopset 17之前要展开成 ReduceMeanSubPow…GridSampleopset 16空间变换常用之前不支持DeformConv至今没有永远需要自定义 domain怎么判断是不是版本问题去 Operators.md查询方式见下文第六节查这个算子的版本历史把opset_version调高试试如 12 → 17看能否导出看报错里是否有 “Unsupported opset” 或 “op X is not supported in opset Y”。在导出的语句中可以通过参数直接配置所需的opset版本号。torch.onnx.export(modelmodel,args(input,),ffile,input_names[input0],output_names[output0],opset_version12)注意不是 opset 越高越好——你的 TensorRT 版本有支持上限例如 TRT 8.5 最高支持 opset 17且高版本可能引入新的节点展开方式。实际工程最常用11 / 12 / 13 / 17。三、4 种注册写法对应不同场景方式适用工况核心代码导出后算子类型register_custom_op_symbolic工况 1词典缺条目register_custom_op_symbolic(aten::xxx, fn, opset)ONNX 标准算子_onnx_symbolic装饰器工况 1PyTorch 源码风格_onnx_symbolic(aten::xxx)ONNX 标准算子自定义 domain 注册工况 2ONNX 没这词g.op(custom::xxx, ...)自定义算子后续需 Pluginautograd.Function.symbolic工况 3自研算子类内定义symbolic静态方法自定义算子后续需 Plugin而遇到工况4直接在导出语句中修改opset_version的值就行。四、关键辨析新手必踩的 3 个坑4.1 两个算子名不是一回事这是最容易混淆的地方register_custom_op_symbolic(aten::asinh,asinh_symbolic,12)# ↑ 第一个参数PyTorch/ATen 算子名defasinh_symbolic(g,input,*,outNone):returng.op(Asinh,input)# ↑ 这里的第一个参数ONNX 算子名register_custom_op_symbolic的第一个参数是 PyTorch 侧的算子名如aten::asinh、torchvision::deform_conv2dg.op()的第一个参数是 ONNX 侧的算子名如Asinh、custom::deform_conv2d。两者不是同一个东西注册算子就是建立这两者之间的映射。4.2 导出阶段 vs 构建阶段“上半场与下半场”这个坑最致命“注册不等于能跑”。注册 symbolic解决的是导出阶段的问题让 ONNX 图能成功生成写 TensorRT Plugin解决的是构建阶段的问题让custom::节点在 TensorRT 里能算出来。如果你映射到标准算子如Asinh导出即终结TensorRT 原生认识如果你映射到custom domain如custom::deform_conv2d导出只是上半场后面必须写 Plugin 才算完。判断顺序能用标准算子表达 → 尽量映射标准算子一步到位实在不行才用 custom domain等于承诺后面要写 Plugin。4.3parse_args的类型字母表parse_args(v,v,v,v,v,i,i,i,i,i,i,i,i,none)里的字母是官方定义的类型描述符按位置一一对应 symbolic 函数的每个参数字母含义转换结果vvalue张量不转换保持图里的张量torch._C.Valueiint转成 Python intislist of int转成 int 列表ffloat转成 Python floatfslist of float转成 float 列表bbool转成 Python boolsstring转成 Python strttorch.Tensor转成张量字面量none参数不用直接忽略它做的事情本质上是把导出器tracer传进来的每个torch._C.Value按你声明的类型解包成 Python 里的普通类型这样你就能在 symbolic 函数里直接当普通值用比如把整数参数读出来当属性。五、参数表的来源模仿官方格式很多同学看dcn_symbolic那一长串参数很懵。它不是凭空编的而是照抄算子底层注册的真实签名。5.1 底层 C 签名的具体实现以 DCN 为例torchvision 源码torchvision/csrc/ops/deform_conv2d.cpp里注册的 ATen 算子是deform_conv2d(Tensor input,Tensor weight,Tensor offset,Tensor mask,Tensor bias,SymInt stride_h,SymInt stride_w,SymInt pad_h,SymInt pad_w,SymInt dilation_h,SymInt dilation_w,SymInt groups,SymInt offset_groups,booluse_mask)-Tensor14 个参数5 个 Tensor 8 个整数 1 个 bool和dcn_symbolic的参数表逐位对应只是变量名稍作改写dilation→dil、groups→n_weight_grps。为什么必须对齐因为导出器遇到torchvision::deform_conv2d节点时会把这个节点的实参按位置交给你的 symbolic 函数。你的函数签名少一个参数、多一个参数、或类型标错导出时就会解析错或直接报错parse_args里甚至写了descriptor 数量少于实参数量会直接抛AssertionError。C 源码在哪pip 安装的包里没有 C 源码csrc/目录不存在只有编译后的_C.so。要看源码去在线仓库登不上 GitHub 用 Gitee 镜像https://gitee.com/mirrors/vision路径torchvision/csrc/ops/deform_conv2d.cpp。5.2 陌生算子的签名查询ONNX 导出时追踪的不是我们写的 Python 函数而是底层 ATen 算子。torchvision.ops.deform_conv2d这种是 Python 包装函数8 个参数含元组它内部会把元组拆成标量、从张量 shape 推断额外参数再调用底层 ATen 算子torchvision::deform_conv2d14 个标量参数。symbolic 函数注册的是底层 ATen 算子所以必须按底层签名写。5.2.1torchvision算子通过算子名称直接拿底层签名最推荐的方法 直接拿底层签名importinspectimporttorchvision.opsprint(inspect.getsource(torchvision.ops.deform_conv2d))翻到源码最后看return torch.ops.xxx.算子名(...)那一行括号里的参数列表就是底层算子的完整签名returntorch.ops.torchvision.deform_conv2d(input,# 1weight,# 2offset,# 3mask,# 4bias,# 5stride_h,# 6 ← stride(h,w) 拆成两个标量stride_w,# 7pad_h,# 8 ← padding(h,w) 拆成两个标量pad_w,# 9dil_h,# 10 ← dilation(h,w) 拆成两个标量dil_w,# 11n_weight_grps,# 12 ← 从 weight.shape 推断n_in_channels // weight.shape[1]n_offset_grps,# 13 ← 从 offset.shape 推断offset.shape[1] // (2 * kh * kw)use_mask,# 14 ← 从 mask is not None 推断的布尔值)我自己的运行结果注意_schemas里的参数名dilation_h、groups和 symbolic 函数里的参数名dil_h、n_weight_grps不完全一样 —— 这没关系。symbolic 函数的参数名是你自己定义的只要参数顺序和类型对得上就行。parse_args里的类型串按顺序对应每个参数。5.2.2 怎么知道算子在哪个命名空间列出所有算子可以通过代码直接查询所有的ATen 算子和torchvision 算子此代码还支持模糊搜索importtorchdefprint_ops(op_list,per_line3,title):整理打印排序 编号 每行固定个数opssorted(op_list)iftitle:print(f\n{title}共{len(ops)}个)foriinrange(0,len(ops),per_line):batchops[i:iper_line]print( .join(f{ij1:3d}.{op}forj,opinenumerate(batch)))# 1. ATen 算子过滤掉 backward只看前向aten_ops[opforopintorch.ops.aten._dirifnotop.endswith(_backward)andnotop.startswith(_)]print_ops(aten_ops,per_line3,titleATen 前向算子)# 2. torchvision 算子过滤掉 backward 和内部属性tv_ops[opforopintorch.ops.torchvision._dirifnotop.endswith(_backward)andnotop.startswith(_)]print_ops(tv_ops,per_line3,titletorchvision 前向算子)# 3. 模糊搜索只打印含关键词的算子 keywordconv# 改成你想搜的关键词matched[opforopintorch.ops.aten._dirifkeywordinop.lower()]print_ops(matched,per_line3,titlefATen 中含 {keyword} 的算子)运行输出如下有个小bug我本地跑的时候发现有时候输出torchvision前向算子为0问题不大本身torchvision就没几个算子平时多看看就记住了。Ten 算子从哪来.so 库 Dispatcher 机制答案是不来自torch/onnx/目录的任何 Python 文件而是来自编译后的.so库。库文件路径内容libtorch_cpu.so/usr/local/lib/python3.10/dist-packages/torch/lib/libtorch_cpu.so所有 ATen 算子的CPU 实现 算子 schema 注册libtorch_cuda.so/usr/local/lib/python3.10/dist-packages/torch/lib/libtorch_cuda.so所有 ATen 算子的CUDA 实现底层实现与上层调用的关系如下5.3 注册算子时模仿的官方格式映射函数代码里面写dcn_symbolic时参照的格式模板就是PyTorch 官方torch/onnx/symbolic_opset*.py。官方文件里一个典型算子的写法真实代码来自symbolic_opset12.py,实际的参数表则是来自于5.1小节的底层实现。_onnx_symbolic(aten::einsum)symbolic_helper.parse_args(s,v,is)_beartype.beartypedefeinsum(g:jit_utils.GraphContext,equation,tensor_list,pathNone):tensorssymbolic_helper._unpack_list(tensor_list)returng.op(Einsum,*tensors,equation_sequation)对比文中的dcn_symbolic每个元素都是一一对应模仿的官方格式symbolic_opset*.py本文案例_onnx_symbolic(aten::einsum)register_custom_op_symbolic(torchvision::deform_conv2d, dcn_symbolic, 12)symbolic_helper.parse_args(s,v,is)parse_args(v,v,v,v,v,i,i,i,i,i,i,i,i,none)def einsum(g, equation, tensor_list)def dcn_symbolic(g, input, weight, offset, ...)return g.op(Einsum, ...)return g.op(custom::deform_conv2d, input, offset)这些文件在电脑上就能找到Python 环境site-packages/torch/onnx/目录下有symbolic_opset7.py~symbolic_opset20.py一整套Jetson 上一般在/usr/local/lib/python3.10/dist-packages/torch/onnx/。打开任何一个就是最权威的格式教科书。六、 ONNX 算子文档获取与理解Operators.md 全攻略一般是去 GitHub 查https://github.com/onnx/onnx/blob/main/docs/Operators.md但 GitHub 经常登不上。以下是全套替代方案方式 1Gitee 镜像最直接内容与 GitHub 完全一致https://gitee.com/mirrors/ONNX/blob/main/docs/Operators.mdGitee 官方同步了 ONNX 仓库国内直接访问。每个算子下面都有版本历史表写清楚从哪个 opset 开始支持、每个版本改了什么。方式 2ONNX 官网带搜索体验更好https://onnx.ai/onnx/operators/有搜索框、每个算子有示例代码和版本切换比 GitHub 的 markdown 好读。偶尔会慢但通常比 GitHub 好连。方式 3本地源码仓库如果有 onnx 源码如果可以将整个onnx源码下载到本地。然后通过VScode打开使用非常方便。Operator列是各个算子Since version列意思是支持变更的历史版本以第一个Abs算子为例显示13,6,1意思是从onnx第1个版本就开始支持了Abs算子后面分别得第6第13个版本进行了小幅改动。点击任意一个版本可以跳转到变更Changelog文档可以看看这几个版本到底有什么变化还是以Abs算子为例。版本1是正常的Abs算子说明支持3种的浮点类型float是32位浮点float16是16位浮点数double是64位浮点数。这里的位指的是浮点数的存储位数具体的可以参考嵌入式视角下的浮点运算性能之迷版本6不仅支持浮点运算也开始支持定点运算。同时取消了consumed_inputs参数。版本13支持的数据类型更多新增bfloat16浮点运算。七Netron 可视化看已有模型用https://netron.app/国内可直接访问。把.onnx文件拖进去点任何一个算子节点右边面板会显示它的op_type、domain、所有属性值、输入输出 shape。调试模型时最直观。八、遇到导出报错的标准排查流程实战总结看报错里是哪个算子如aten::asinh、torchvision::deform_conv2d查 ONNX 标准里有没有它用第六节的任一方式有→ 工况 1注册映射到标准算子搞定没有→ 工况 2注册自定义 domain然后准备写 Plugin或先用 onnx-graphsurgeon 拆图是你自己写的 Function→ 工况 3决定要不要定义 symbolicopset 报错→ 工况 4先试 12 / 13不行再往高调。九、总结把全文串成一条主线PyTorch 模型里出现了一个算子不认识 │ ├─ ONNX 标准里有它─────────是──▶ 注册映射成标准算子如 asinh→ 到此为止 ✅ │ └─ 没有 ──▶ 注册成 custom:: 节点DCN / 自研算子→ 只是上半场 │ ├─ 能用标准算子组合表达──是──▶ onnx-graphsurgeon 拆图改 ONNX 文件 │ └─ 拆不了 ──▶ 写 TensorRT Plugin 补底层C/CUDA→ 下半场 ✅一句话记忆注册算子就是给PyTorch→ONNX这本翻译词典补条目。有血缘ONNX 标准里有介绍一句就认识没血缘ONNX 标准里没有就得自己造底层证明Plugin。查算子文档不用死磕 GitHubGitee 镜像、onnx.ai、本地 onnx 包、Netron 随便选。希望这篇文章能帮你打通算子注册这一关。下一篇可以聊聊 onnx-graphsurgeon 改图拆图 vs 写 Plugin 的抉择欢迎关注。参考资料PyTorch 官方torch/onnx/symbolic_opset*.py、torchvision 源码torchvision/csrc/ops/deform_conv2d.cpp、ONNX 官方文档 优雅地解决 onnx 算子不兼容 —— 算子注册
目录前言为什么要看懂算子注册一、ONNX 导出本质上是翻译二、4 种算子不支持的工况工况 1词典缺条目PyTorch 有、ONNX 标准也有、只是导出器没映射工况 2目标语言没这个词torchvision 有、ONNX 标准没有目标语言没这个词torchvision 有、ONNX 标准没有工况 3自研算子autograd.Function工况 4opset 版本问题用了旧版词典三、4 种注册写法对应不同场景四、关键辨析新手必踩的 3 个坑4.1 两个算子名不是一回事4.2 导出阶段 vs 构建阶段上半场与下半场4.3 parse_args 的类型字母表五、参数表的来源模仿官方格式5.1 底层 C 签名的具体实现5.2 陌生算子的签名查询5.2.1 **torchvision算子通过算子名称直接拿底层签名**5.2.2 怎么知道算子在哪个命名空间列出所有算子5.3 注册算子时模仿的官方格式映射函数六、 ONNX 算子文档获取与理解Operators.md 全攻略方式 1Gitee 镜像最直接内容与 GitHub 完全一致方式 2ONNX 官网带搜索体验更好方式 3本地源码仓库如果有 onnx 源码七Netron 可视化看已有模型用八、遇到导出报错的标准排查流程实战总结九、总结前言为什么要看懂算子注册做模型部署的人几乎每天都在跟这条流水线打交道而**算子不支持是这条流水线上最高频的拦路虎**。你在 PyTorch 里写得好好的一个模型一执行torch.onnx.export就报一堆看不懂的错好不容易导出成功了喂给trtexec构建 engine 又报错。本文以及后续几篇文章要解决的核心问题就一个当算子不认识时到底该怎么办本文将给出 4 种工况针对不兼容的4种工况给出相应的解决方案最后附上去哪里查算子文档的保姆级清单。一、ONNX 导出本质上是翻译PyTorch 导出的过程本质上是一个翻译过程PyTorch 的计算图是用ATen 算子a Tensor LibraryPyTorch 底层的 C 张量运算库写成的——可以理解为中文ONNX 的计算图是用ONNX 标准算子写成的——可以理解为英文导出器torch.onnx.export就是一个翻译官它手里有一本词典记录了每个 ATen 算子中文词对应哪个 ONNX 算子英文词。这本词典就是 PyTorch 源码里的torch/onnx/symbolic_opset*.py文件。每一个词条叫符号函数symbolic function。如果词典里没有某个中文词的条目翻译官就卡壳了于是导出报错。注册算子就是手动给这本词典补一条新条目“这个中文词aten::xxx翻译成英文词Xxx这样翻。”二、4 种算子不支持的工况把报错归归类你会发现算子不支持其实分 4 种完全不同的情况处理方式天差地别。工况 1词典缺条目PyTorch 有、ONNX 标准也有、只是导出器没映射代表torch.asinh示例工程 3.4 的sample_asinh.py现象直接导出报错。根因PyTorch 有torch.asinh()ATen 里有aten::asinh✅ONNX 标准从opset 9起就有Asinh算子 ✅但是用的 PyTorch 版本的symbolic_opset*.py里没写aten::asinh→Asinh的映射条目 ❌就像中文有asinh这个词、英文也有Asinh这个词但翻译官的词典漏了这一条。解决fromtorch.onnximportregister_custom_op_symbolicdefasinh_symbolic(g,input,*,outNone):returng.op(Asinh,input)# 左边是 ONNX 算子名register_custom_op_symbolic(aten::asinh,asinh_symbolic,12)# ↑ PyTorch 算子名就 3 行核心代码手动补一条翻译条目。导出时遇到aten::asinh就知道要生成 ONNX 的Asinh节点。工况 2目标语言没这个词torchvision 有、ONNX 标准没有目标语言没这个词torchvision 有、ONNX 标准没有代表torchvision.ops.DeformConv2d现象直接导出报错。根因torchvision 有DeformConv2d底层有完整 CUDA 实现 ✅ONNX 标准里从来没有可变形卷积这个算子❌就像中文有内卷这个词但英文里根本没有对应的标准词汇——你不能把它翻译成一个不存在的词。解决sample_deformable_conv_register.pyfromtorch.onnximportregister_custom_op_symbolicfromtorch.onnx.symbolic_helperimportparse_argsparse_args(v,v,v,v,v,i,i,i,i,i,i,i,i,none)defdcn_symbolic(g,input,weight,offset,mask,bias,stride_h,stride_w,pad_h,pad_w,dil_h,dil_w,n_weight_grps,n_offset_grps,use_mask):returng.op(custom::deform_conv2d,input,offset)# 注意自定义 domainregister_custom_op_symbolic(torchvision::deform_conv2d,dcn_symbolic,12)注意这里的关键区别g.op(custom::deform_conv2d, ...)用了自定义 domaincustom::而不是g.op(DeformConv2d, ...)——因为 ONNX 标准里没这个算子你不能凭空编一个标准算子。这意味着什么导出成custom::deform_conv2d后问题并没有真正解决完只是上半场onnxruntime 默认跑不了不认识这个算子TensorRT 也跑不了必须自己写一个TensorRT PluginC/CUDA kernel来实现可变形卷积的前向计算也就是说我们虽然可以顺利导出onnx文件但实质上这部分是个“空壳”这就是下半场后面有时间再更新相关的文章敬请期待。提醒可变形卷积DCN在目标检测模型CenterPoint、FCOS3D 等中非常常见部署时几乎都要写 TensorRT Plugin。这是部署工程师的核心技能之一。工况 3自研算子autograd.Function代表自己写的一个torch.autograd.Function算子。现象这一组不是报错而是导出结果不符合预期——你写的自定义算子被**拆开inline**成了一串基础节点。classCustomOp(torch.autograd.Function):staticmethoddefforward(ctx,x:torch.Tensor)-torch.Tensor:ctx.save_for_backward(x)xx.clamp(min0)returnx/(1torch.exp(-x))导出后的onnx文件网络结构是这样的根因如果用torch.autograd.Function写了一个自定义算子导出器在不注册 symbolic时会追踪 forward 里的每一个基础算子最终导出的 ONNX 图是Clip → Exp → Add → Div一串节点。图能跑但不是一个整体算子。解决在 Function 子类里定义symbolic静态方法importtorchclassCustomOp(torch.autograd.Function):staticmethoddefsymbolic(g,x):returng.op(custom_domain::customOp2,x)# 导出成单个自定义节点staticmethoddefforward(ctx,x):xx.clamp(min0)returnx/(1torch.exp(-x))customOpCustomOp.apply导出后的onnx文件网络结构是这样的什么时候用哪种如果自定义算子的计算能用 ONNX 标准算子组合表达、且性能没问题 → 不注册 symbolic让它 inline 展开就行如果需要在推理引擎里对算子做特殊优化融合、写 Plugin→ 必须注册 symbolic 导出为单个节点否则引擎看不到这是一个整体没法优化。工况 4opset 版本问题用了旧版词典现象同一个算子换个 opset 版本就能导出或换个版本就报错。根因ONNX 标准在演进每个 opset 版本会新增 / 废弃 / 修改算子。常见例子算子引入版本说明Asinhopset 9更低版本就没有LayerNormalizationopset 17之前要展开成 ReduceMeanSubPow…GridSampleopset 16空间变换常用之前不支持DeformConv至今没有永远需要自定义 domain怎么判断是不是版本问题去 Operators.md查询方式见下文第六节查这个算子的版本历史把opset_version调高试试如 12 → 17看能否导出看报错里是否有 “Unsupported opset” 或 “op X is not supported in opset Y”。在导出的语句中可以通过参数直接配置所需的opset版本号。torch.onnx.export(modelmodel,args(input,),ffile,input_names[input0],output_names[output0],opset_version12)注意不是 opset 越高越好——你的 TensorRT 版本有支持上限例如 TRT 8.5 最高支持 opset 17且高版本可能引入新的节点展开方式。实际工程最常用11 / 12 / 13 / 17。三、4 种注册写法对应不同场景方式适用工况核心代码导出后算子类型register_custom_op_symbolic工况 1词典缺条目register_custom_op_symbolic(aten::xxx, fn, opset)ONNX 标准算子_onnx_symbolic装饰器工况 1PyTorch 源码风格_onnx_symbolic(aten::xxx)ONNX 标准算子自定义 domain 注册工况 2ONNX 没这词g.op(custom::xxx, ...)自定义算子后续需 Pluginautograd.Function.symbolic工况 3自研算子类内定义symbolic静态方法自定义算子后续需 Plugin而遇到工况4直接在导出语句中修改opset_version的值就行。四、关键辨析新手必踩的 3 个坑4.1 两个算子名不是一回事这是最容易混淆的地方register_custom_op_symbolic(aten::asinh,asinh_symbolic,12)# ↑ 第一个参数PyTorch/ATen 算子名defasinh_symbolic(g,input,*,outNone):returng.op(Asinh,input)# ↑ 这里的第一个参数ONNX 算子名register_custom_op_symbolic的第一个参数是 PyTorch 侧的算子名如aten::asinh、torchvision::deform_conv2dg.op()的第一个参数是 ONNX 侧的算子名如Asinh、custom::deform_conv2d。两者不是同一个东西注册算子就是建立这两者之间的映射。4.2 导出阶段 vs 构建阶段“上半场与下半场”这个坑最致命“注册不等于能跑”。注册 symbolic解决的是导出阶段的问题让 ONNX 图能成功生成写 TensorRT Plugin解决的是构建阶段的问题让custom::节点在 TensorRT 里能算出来。如果你映射到标准算子如Asinh导出即终结TensorRT 原生认识如果你映射到custom domain如custom::deform_conv2d导出只是上半场后面必须写 Plugin 才算完。判断顺序能用标准算子表达 → 尽量映射标准算子一步到位实在不行才用 custom domain等于承诺后面要写 Plugin。4.3parse_args的类型字母表parse_args(v,v,v,v,v,i,i,i,i,i,i,i,i,none)里的字母是官方定义的类型描述符按位置一一对应 symbolic 函数的每个参数字母含义转换结果vvalue张量不转换保持图里的张量torch._C.Valueiint转成 Python intislist of int转成 int 列表ffloat转成 Python floatfslist of float转成 float 列表bbool转成 Python boolsstring转成 Python strttorch.Tensor转成张量字面量none参数不用直接忽略它做的事情本质上是把导出器tracer传进来的每个torch._C.Value按你声明的类型解包成 Python 里的普通类型这样你就能在 symbolic 函数里直接当普通值用比如把整数参数读出来当属性。五、参数表的来源模仿官方格式很多同学看dcn_symbolic那一长串参数很懵。它不是凭空编的而是照抄算子底层注册的真实签名。5.1 底层 C 签名的具体实现以 DCN 为例torchvision 源码torchvision/csrc/ops/deform_conv2d.cpp里注册的 ATen 算子是deform_conv2d(Tensor input,Tensor weight,Tensor offset,Tensor mask,Tensor bias,SymInt stride_h,SymInt stride_w,SymInt pad_h,SymInt pad_w,SymInt dilation_h,SymInt dilation_w,SymInt groups,SymInt offset_groups,booluse_mask)-Tensor14 个参数5 个 Tensor 8 个整数 1 个 bool和dcn_symbolic的参数表逐位对应只是变量名稍作改写dilation→dil、groups→n_weight_grps。为什么必须对齐因为导出器遇到torchvision::deform_conv2d节点时会把这个节点的实参按位置交给你的 symbolic 函数。你的函数签名少一个参数、多一个参数、或类型标错导出时就会解析错或直接报错parse_args里甚至写了descriptor 数量少于实参数量会直接抛AssertionError。C 源码在哪pip 安装的包里没有 C 源码csrc/目录不存在只有编译后的_C.so。要看源码去在线仓库登不上 GitHub 用 Gitee 镜像https://gitee.com/mirrors/vision路径torchvision/csrc/ops/deform_conv2d.cpp。5.2 陌生算子的签名查询ONNX 导出时追踪的不是我们写的 Python 函数而是底层 ATen 算子。torchvision.ops.deform_conv2d这种是 Python 包装函数8 个参数含元组它内部会把元组拆成标量、从张量 shape 推断额外参数再调用底层 ATen 算子torchvision::deform_conv2d14 个标量参数。symbolic 函数注册的是底层 ATen 算子所以必须按底层签名写。5.2.1torchvision算子通过算子名称直接拿底层签名最推荐的方法 直接拿底层签名importinspectimporttorchvision.opsprint(inspect.getsource(torchvision.ops.deform_conv2d))翻到源码最后看return torch.ops.xxx.算子名(...)那一行括号里的参数列表就是底层算子的完整签名returntorch.ops.torchvision.deform_conv2d(input,# 1weight,# 2offset,# 3mask,# 4bias,# 5stride_h,# 6 ← stride(h,w) 拆成两个标量stride_w,# 7pad_h,# 8 ← padding(h,w) 拆成两个标量pad_w,# 9dil_h,# 10 ← dilation(h,w) 拆成两个标量dil_w,# 11n_weight_grps,# 12 ← 从 weight.shape 推断n_in_channels // weight.shape[1]n_offset_grps,# 13 ← 从 offset.shape 推断offset.shape[1] // (2 * kh * kw)use_mask,# 14 ← 从 mask is not None 推断的布尔值)我自己的运行结果注意_schemas里的参数名dilation_h、groups和 symbolic 函数里的参数名dil_h、n_weight_grps不完全一样 —— 这没关系。symbolic 函数的参数名是你自己定义的只要参数顺序和类型对得上就行。parse_args里的类型串按顺序对应每个参数。5.2.2 怎么知道算子在哪个命名空间列出所有算子可以通过代码直接查询所有的ATen 算子和torchvision 算子此代码还支持模糊搜索importtorchdefprint_ops(op_list,per_line3,title):整理打印排序 编号 每行固定个数opssorted(op_list)iftitle:print(f\n{title}共{len(ops)}个)foriinrange(0,len(ops),per_line):batchops[i:iper_line]print( .join(f{ij1:3d}.{op}forj,opinenumerate(batch)))# 1. ATen 算子过滤掉 backward只看前向aten_ops[opforopintorch.ops.aten._dirifnotop.endswith(_backward)andnotop.startswith(_)]print_ops(aten_ops,per_line3,titleATen 前向算子)# 2. torchvision 算子过滤掉 backward 和内部属性tv_ops[opforopintorch.ops.torchvision._dirifnotop.endswith(_backward)andnotop.startswith(_)]print_ops(tv_ops,per_line3,titletorchvision 前向算子)# 3. 模糊搜索只打印含关键词的算子 keywordconv# 改成你想搜的关键词matched[opforopintorch.ops.aten._dirifkeywordinop.lower()]print_ops(matched,per_line3,titlefATen 中含 {keyword} 的算子)运行输出如下有个小bug我本地跑的时候发现有时候输出torchvision前向算子为0问题不大本身torchvision就没几个算子平时多看看就记住了。Ten 算子从哪来.so 库 Dispatcher 机制答案是不来自torch/onnx/目录的任何 Python 文件而是来自编译后的.so库。库文件路径内容libtorch_cpu.so/usr/local/lib/python3.10/dist-packages/torch/lib/libtorch_cpu.so所有 ATen 算子的CPU 实现 算子 schema 注册libtorch_cuda.so/usr/local/lib/python3.10/dist-packages/torch/lib/libtorch_cuda.so所有 ATen 算子的CUDA 实现底层实现与上层调用的关系如下5.3 注册算子时模仿的官方格式映射函数代码里面写dcn_symbolic时参照的格式模板就是PyTorch 官方torch/onnx/symbolic_opset*.py。官方文件里一个典型算子的写法真实代码来自symbolic_opset12.py,实际的参数表则是来自于5.1小节的底层实现。_onnx_symbolic(aten::einsum)symbolic_helper.parse_args(s,v,is)_beartype.beartypedefeinsum(g:jit_utils.GraphContext,equation,tensor_list,pathNone):tensorssymbolic_helper._unpack_list(tensor_list)returng.op(Einsum,*tensors,equation_sequation)对比文中的dcn_symbolic每个元素都是一一对应模仿的官方格式symbolic_opset*.py本文案例_onnx_symbolic(aten::einsum)register_custom_op_symbolic(torchvision::deform_conv2d, dcn_symbolic, 12)symbolic_helper.parse_args(s,v,is)parse_args(v,v,v,v,v,i,i,i,i,i,i,i,i,none)def einsum(g, equation, tensor_list)def dcn_symbolic(g, input, weight, offset, ...)return g.op(Einsum, ...)return g.op(custom::deform_conv2d, input, offset)这些文件在电脑上就能找到Python 环境site-packages/torch/onnx/目录下有symbolic_opset7.py~symbolic_opset20.py一整套Jetson 上一般在/usr/local/lib/python3.10/dist-packages/torch/onnx/。打开任何一个就是最权威的格式教科书。六、 ONNX 算子文档获取与理解Operators.md 全攻略一般是去 GitHub 查https://github.com/onnx/onnx/blob/main/docs/Operators.md但 GitHub 经常登不上。以下是全套替代方案方式 1Gitee 镜像最直接内容与 GitHub 完全一致https://gitee.com/mirrors/ONNX/blob/main/docs/Operators.mdGitee 官方同步了 ONNX 仓库国内直接访问。每个算子下面都有版本历史表写清楚从哪个 opset 开始支持、每个版本改了什么。方式 2ONNX 官网带搜索体验更好https://onnx.ai/onnx/operators/有搜索框、每个算子有示例代码和版本切换比 GitHub 的 markdown 好读。偶尔会慢但通常比 GitHub 好连。方式 3本地源码仓库如果有 onnx 源码如果可以将整个onnx源码下载到本地。然后通过VScode打开使用非常方便。Operator列是各个算子Since version列意思是支持变更的历史版本以第一个Abs算子为例显示13,6,1意思是从onnx第1个版本就开始支持了Abs算子后面分别得第6第13个版本进行了小幅改动。点击任意一个版本可以跳转到变更Changelog文档可以看看这几个版本到底有什么变化还是以Abs算子为例。版本1是正常的Abs算子说明支持3种的浮点类型float是32位浮点float16是16位浮点数double是64位浮点数。这里的位指的是浮点数的存储位数具体的可以参考嵌入式视角下的浮点运算性能之迷版本6不仅支持浮点运算也开始支持定点运算。同时取消了consumed_inputs参数。版本13支持的数据类型更多新增bfloat16浮点运算。七Netron 可视化看已有模型用https://netron.app/国内可直接访问。把.onnx文件拖进去点任何一个算子节点右边面板会显示它的op_type、domain、所有属性值、输入输出 shape。调试模型时最直观。八、遇到导出报错的标准排查流程实战总结看报错里是哪个算子如aten::asinh、torchvision::deform_conv2d查 ONNX 标准里有没有它用第六节的任一方式有→ 工况 1注册映射到标准算子搞定没有→ 工况 2注册自定义 domain然后准备写 Plugin或先用 onnx-graphsurgeon 拆图是你自己写的 Function→ 工况 3决定要不要定义 symbolicopset 报错→ 工况 4先试 12 / 13不行再往高调。九、总结把全文串成一条主线PyTorch 模型里出现了一个算子不认识 │ ├─ ONNX 标准里有它─────────是──▶ 注册映射成标准算子如 asinh→ 到此为止 ✅ │ └─ 没有 ──▶ 注册成 custom:: 节点DCN / 自研算子→ 只是上半场 │ ├─ 能用标准算子组合表达──是──▶ onnx-graphsurgeon 拆图改 ONNX 文件 │ └─ 拆不了 ──▶ 写 TensorRT Plugin 补底层C/CUDA→ 下半场 ✅一句话记忆注册算子就是给PyTorch→ONNX这本翻译词典补条目。有血缘ONNX 标准里有介绍一句就认识没血缘ONNX 标准里没有就得自己造底层证明Plugin。查算子文档不用死磕 GitHubGitee 镜像、onnx.ai、本地 onnx 包、Netron 随便选。希望这篇文章能帮你打通算子注册这一关。下一篇可以聊聊 onnx-graphsurgeon 改图拆图 vs 写 Plugin 的抉择欢迎关注。参考资料PyTorch 官方torch/onnx/symbolic_opset*.py、torchvision 源码torchvision/csrc/ops/deform_conv2d.cpp、ONNX 官方文档