pytorch自定义算子导出onnx

1、为什么要自定义算子？

1、没有现成可用的算子，需要根据自己的接口重写。
2、现有的算子接口不兼容，需要在原有的算子上进行封装。

2、如何自定义算子

继承torch.autograd.Function类，实现其forward() 和 backward()方法,就可以成为一个pytorch自定义算子。就可以在模型训练推理中完成前向推理和反向传播。
forward() 函数的第一个参数必须是ctx, 后面是输入。
在工程部署上，一般为了加快计算，自定义算子需要用cuda 实现forward()、backward（）kernel 函数。

3、自定义算子导出onnx

实现其symbolic 静态方法，当我们调用torch.onnx.export()时，就可以导出onnx 算子。
symbolic是符号函数，通常在其内部返回一个g.op()对象。g.op() 把一个 PyTorch 算子映射成一个或多个 ONNX 算子，或者是自定义的 ONNX 算子。
symbolic函数的第一个参数必须是g, 后面是和forward()对应的输入。
g.op() 做算子映射，g.op 的参数：
1、第一个参数为算子名字，
2、后面参数与forward() 输入对应，
3、往后可以是一些算子自带常量和属性值。常量视为输入，属性值需要用 字段_s/i/f = 默认值表示。_s 表示字符串，_i 表示 int64, _f 表示 float32。常量用类似 g.op(“Constant”, value_t=torch.tensor([3, 2, 1], dtype=torch.float32))表示

4、example

1、重写一个pytorch 自定义算子（实现自定义激活函数）

实现自己的激活函数MYSELU 算子。

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.onnx
import torch.autograd

#继承torch.autograd.Function
class MYSELUImpl(torch.autograd.Function): 
    @staticmethod
    def symbolic(g, x, p):
        return g.op("MYSELU", x, p,  # 表示onnx算子的名称为MYSELU，参数与forward()对应
            # 给算子传一个常数参数
            g.op("Constant", value_t=torch.tensor([3, 2, 1], dtype=torch.float32)),
            attr1_s="这是字符串属性", # s表示字符串
            attr2_i=[1, 2, 3], # i表示整数
            attr3_f=222  # f表示浮点数
        )
    @staticmethod
    def forward(ctx, x