pytorch自定义算子并导出onnx计算图详细代码教程

  官方教程和其他博客只对极其简单的模型进行了介绍，由于使用中会涉及到一些教程中没有撰写的东西，因此写下本博客进行记录，直接上代码，最后进行总结避坑。

单输入多输出模型

代码

import torch
import torch.nn as nn
from torch.autograd import Function
import torch.onnx

#custom op onnx representation
class MyStrangeOp(Function):
   #for onnx node representation, symbolic function must be defined and specified static.
   @staticmethod
   def symbolic(g, input, weight, bias, floatAttr, intAttr):
      #because forward function return 2 outputs, so this func also have to return 2 outputs
      #this is my expriment result, I didn't find any docs, fuck it!
      return g.op("MyStrangeOp", input, weight, bias, float_attr_f=floatAttr, int_attr_i=intAttr), \
             g.op("MyStrangeOp", input, weight, bias, float_attr_f=floatAttr, int_attr_i=intAttr)

   @staticmethod
   def forward(ctx, input, weight, bias, floatAttr, intAttr):
      #this op return 2 outputs
      return input + weight, input * weight + bias

myStrangeOpForward = MyStrangeOp.apply

#layer
class MyStrangeOpLayer(nn.Module):
   def __init__(self, weight, bias, floatAttr, intAttr):
      super(MyStrangeOpLayer, self).__init__()
      self.weight = weight
      self.bias = bias
      self.floatAttr = floatAttr
      self.intAttr = intAttr

   def forward(self, x):
      return myStrangeOpForward(x, self.weight, self.bias, self.floatAttr, self.intAttr)

#model
class MyStrangeNet(nn.Module):
   def __init__(self):
      super(MyStrangeNet, self).__init__()
      self.myLayer1 = MyStrangeOpLayer(weight=nn.Parameter(torch.ones(1, 3, 4, 4)), bias=nn.Parameter(torch.ones(1, 3, 4, 4)), floatAttr=[0.1, 0.5], intAttr=[2, 2])
      self.myLayer2 = MyStrangeOpLayer(weight=nn.Parameter(torch.ones(1, 3, 4, 4)), bias=nn.Parameter(torch.ones(1, 3, 4, 4)), floatAttr=[0.5, 0.5], intAttr=[3, 3])
      self.conv1    = nn.Conv2d(in_channels=3, out_channels=3, kernel_size=3, padding=1, stride=1, bias=True)
      self.conv2    = nn.Conv2d(in_channels=3, out_channels=3, kernel_size=3, padding=1, stride=1, bias=True)
      self.conv3    = nn.Conv2d(in_channels=3, out_channels=3, kernel_size=3, padding=1, stride=1, bias=True)
      self.conv4    = nn.Conv2d(in_channels=3, out_channels=3, kernel_size=3, padding=1, stride=1, bias=True)

   def forward(self, x):
      x1, x2 = self.myLayer1(x)
      x3, x4 = self.myLayer2(x)
      x1     = self.conv1(x1)
      x2     = self.conv1(x2)
      x3     = self.conv1(x3)
      x4     = self.conv1(x4)
      return x1 + x2 + x3 + x4



model = MyStrangeNet()
t = torch.ones(1, 3, 4, 4, dtype=torch.float32)
torch.onnx.export(model, (t,), 'fuckIt.onnx', opset_version=13, input_names=["inputTensor"], output_names=["outputTensor"],
                  operator_export_type=torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK)

模型onnx图

本模型经验总结

1. 多输出模型在定义symbolic函数时，需要使得返回值个数等于输出个数，返回值一直重复第一个返回值的内容即可，我没有在任何博客或官方文档中看到多返回值的处理情况，这完全是通过报错然后瞎改发现的。

多输入单输出模型

代码

import torch
import torch.nn as nn
from torch.autograd import Function
import torch.onnx

#custom op for onnx representation
class MyStrangeOp2(Function):
	#for onnx graph
   @staticmethod
   def symbolic(g, input1, input2, bias, int_attr1, int_attr2, str_attr3):
      return g.op("MyStrangeOp2", input1, input2, bias, int_attr1_i=int_attr1, int_attr2_i=int_attr2, str_attr3_s=str_attr3)

   @staticmethod
   def forward(ctx, input1, input2, bias, int_attr1, int_attr2, str_attr3):
      return input1 + input2 - bias

myStrangeOp2_forward = MyStrangeOp2.apply

#layer
class MyStrangeOp2Layer(nn.Module):
   def __init__(self, bias, int_attr1, int_attr2, str_attr3):
	   super(MyStrangeOp2Layer, self).__init__()
	   self.bias = bias
	   self.int_attr1 = int_attr1
	   self.int_attr2 = int_attr2
	   self.str_attr3 = str_attr3

   def forward(self, in1, in2):
	   return myStrangeOp2_forward(in1, in2, self.bias, self.int_attr1, self.int_attr2, self.str_attr3)

#net
class MyStrangeNet2(nn.Module):
   def __init__(self):
      super(MyStrangeNet2, self).__init__()
      self.myLayer1 = MyStrangeOp2Layer(bias=nn.Parameter(torch.ones(1, 3, 4, 4)), int_attr1=10, int_attr2=[3, 5], str_attr3="fuck" )
      self.myLayer2 = MyStrangeOp2Layer(bias=nn.Parameter(torch.ones(1, 3, 4, 4)), int_attr1=40, int_attr2=[12, 22], str_attr3="shit")
      self.conv1    = nn.Conv2d(in_channels=3, out_channels=3, kernel_size=3, padding=1, stride=1, bias=True)
      self.conv2    = nn.Conv2d(in_channels=3, out_channels=3, kernel_size=3, padding=1, stride=1, bias=True)

   def forward(self, in1, in2, in3, in4):
      x1 = self.myLayer1(in1, in2)
      x2 = self.myLayer2(in3, in4)
      x1 = self.conv1(x1)
      x2 = self.conv2(x2)
      return x1 + x2


#fake input
model = MyStrangeNet2()
t1 = torch.ones(1, 3, 4, 4, dtype=torch.float32)
t2= torch.ones(1, 3, 4, 4, dtype=torch.float32)
t3 = torch.ones(1, 3, 4, 4, dtype=torch.float32)
t4 = torch.ones(1, 3, 4, 4, dtype=torch.float32)
#save onnx
torch.onnx.export(model, (t1, t2, t3, t4), 'fuckNet.onnx', opset_version=13, input_names=["input1", "input2", "input3", "input4",], output_names=["outputTensor"],
                  operator_export_type=torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK)

模型onnx图

本模型经验总结

多输入单输出模型这个细分模型没有太多要注意的地方。

结论

除了上述两个模型分别的一些特性之外，所有自定义算子的实现过程都有一些共性特点，可总结如下：

1. 调用torch.onnx.export函数保存onnx模型时，有的版本需要加上operator_export_type=torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK防止报错，有的版本不需要。
2. 定义一个自定义节点时，在symbolic函数中的g.op(xxx)函数内，若输入实参不带关键字，则onnx节点节点会将该数据解析为input，如果输入的实参带关键字，节点会将该数据解析为attribute。对于input又可以分为确实是需要的输入或是网络层的权重属性，后者如果是nn.Parameter类，则在onnx内被分类为initializer，否则为constant。当数据被解析为attribute时，根据关键字的后缀不同(s字符串、f浮点数、i整形)会分配不同的类型，如果他们是一个python列表有多个元素，则在onnx节点内该属性会被解析为列表属性。
3. 不需要将自定义节点拿进pytorch框架进行训练时，是不需要定义backward函数的。如果onnx模型不用于拿来计算而仅仅是作为中间模型最终转化为推理引擎的模型，则forward函数也只需要保证输入输出的个数和shape符合要求即可。
4. 对于nn.Conv2d节点，转化出的onnx模型在netron可视化软件中，输入的代号分别为X W B，输出为Y，这个代号自定义节点只能为0 1 2，这个我不知道是否可以改变也没有找到任何相关的资料，但该代号对节点影响不大，可以忽略。

后记

  本博客的实用性相对而言不太高，假设我们是做模型部署工作，得到了onnx模型，有的节点想要替换很可能还需要拿到python网络的源码，然后在pytorch内对整个网络进行修改，包括现有节点的移除以及自定义节点的添加，当我们拥有了英伟达推出的graph_surgeon工具后，这一切都变得如此简单！传送门如下:链接: ONNX GraphSurgeon API 不得不赞一个nvidia，让模型修改如此的简单，后期有空我会再写一篇gs的教学，很多博客只停留在了API的解释上，正确理解计算图、节点、输入输出这些概念，是我们熟练使用gs的关键！