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): @staticmethoddef 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表示浮点数)@staticmethoddef forward(ctx, x, p): # 前行推理return x * 1 / (1 + torch.exp(-x))class MYSELU(nn.Module): def __init__(self, n):super().__init__()self.param = nn.parameter.Parameter(torch.arange(n).float())def forward(self, x):return MYSELUImpl.apply(x, self.param) #推理调用class Model(nn.Module):def __init__(self):super().__init__()self.conv = nn.Conv2d(1, 1, 3, padding=1)self.myselu = MYSELU(3)self.conv.weight.data.fill_(1)self.conv.bias.data.fill_(0)def forward(self, x):x = self.conv(x)x = self.myselu(x)return x

2、现有算子上封装pytorch 自定义算子（实现动态放大超分辨率模型）

实现动态放大超分辨率模型。我们希望实现：
forward(self, x, upscale_factor)
这样一个接口，x 为图像输入，upscale_factor为动态放大倍数。
pytorch 现有放大算子有nn.Upsample 和 interpolate， 但是nn.Upsample 在初始化阶段固化了放大倍数，而 PyTorch 的 interpolate 插值算子可以在运行阶段选择放大倍数。

class SuperResolutionNet(nn.Module): def forward(self, x, upscale_factor): x = interpolate(x, scale_factor=upscale_factor.item(), mode='bicubic', align_corners=False) 
... 
# Inference 
# Note that the second input is torch.tensor(3) 
torch_output = model(torch.from_numpy(input_img), torch.tensor(3)).detach().numpy() 
... 
with torch.no_grad(): torch.onnx.export(model, (x, torch.tensor(3)), "srcnn2.onnx", opset_version=11, input_names=['input', 'factor'], output_names=['output']) 

尝试使用以上方法导出onnx 时，虽然没有报错能成功导出onnx,但是有TraceWarning 的警告，说明导出onnx有追踪失败。这是由于我们使用了 torch.Tensor.item() 把数据从 Tensor 里取出来，而导出 ONNX 模型时这个操作是无法被记录的，只好报了一条 TraceWarning。

因此我们需要自定义算子，让onnx在追踪时刻能work。我们看到nn.Upsample 和 interpolate在转onnx时都映射到了Resize 操作。所以自定义算子在Resize 操作上进行封装即可。

Resize 操作有三个输入，x， roi， scale, 我们就是要动态输入scale。展开 scales，可以看到 scales 是一个长度为 4 的一维张量，其内容为 [1, 1, 3, 3],
如果我们能够自己生成一个 ONNX 的 Resize 算子，让 scales 成为一个可变量而不是常量，就像它上面的 X 一样，那这个超分辨率模型就能动态缩放了。

import torch 
from torch import nn 
from torch.nn.functional import interpolate 
import torch.onnx 
import cv2 
import numpy as np 
class NewInterpolate(torch.autograd.Function): @staticmethod def symbolic(g, input, scales): return g.op("Resize", input, g.op("Constant", value_t=torch.tensor([], dtype=torch.float32)), scales, coordinate_transformation_mode_s="pytorch_half_pixel", cubic_coeff_a_f=-0.75, mode_s='cubic', nearest_mode_s="floor") @staticmethod def forward(ctx, input, scales): scales = scales.tolist()[-2:] return interpolate(input, scale_factor=scales, mode='bicubic', align_corners=False) class StrangeSuperResolutionNet(nn.Module): def __init__(self): super().__init__() self.conv1 = nn.Conv2d(3, 64, kernel_size=9, padding=4) self.conv2 = nn.Conv2d(64, 32, kernel_size=1, padding=0) self.conv3 = nn.Conv2d(32, 3, kernel_size=5, padding=2) self.relu = nn.ReLU() def forward(self, x, upscale_factor): x = NewInterpolate.apply(x, upscale_factor) out = self.relu(self.conv1(x)) out = self.relu(self.conv2(out)) out = self.conv3(out) return out 

以上自定义了Resize 算子，将scale 作为算子的一个输入，最后还是调用interpolate。但是scale已经变成自定义输入参数。
参数映射如下：