[pytorch] torch.onnx.export 使用及自定义层构建

可以在构建pytorch的时候考虑自定义层能直接转onnx，可以实现带自定义层的pytorch —> onnx 一步到位

• 带参数和不带参数的自定义层构建
• 多个输出的自定义层构建
• 修剪onnx graph

1. torch.onnx.export 参数

torch.onnx.export(
model, 
# 要转onnx的pytorch 模型
args, 
# model forward的函数的Tensor类型的输入， 
# 这个输入和g.op("ai.onnx.contrib::customlayer", intput, args_f=23.0, args_i=8) 注册中的input保持一致， 
# 后面*_f或者*_i的表示是属性， 不算作输入的。
f, 
# f(str): *.onnx文件输出 path and name
export_params=True, 
verbose=False, 
# verbose:如果为True的话，会导致onnx转化的节点description很大
training=TrainingMode.EVAL,
input_names=None, 
# input_names: 和输出设置一样
output_names=None, 
# output_names（list of str, default empty list）： 这个只是设定 输出的名字的，
# 可以和模型中的输出顺序保持一直，从而指定一些输出的名字
operator_export_type=None,
opset_version=None,
# opset_version: 这个表示使用哪个版本的ops集合，这个opt支撑哪些层，
# 可以去源码的torch/onnx/symbolic_opset8-15.py里面去看，支持哪些层，版本越高，支持的越好。
do_constant_folding=True, 
dynamic_axes=None,
keep_initializers_as_inputs=None, 
custom_opsets=None, 
# custom_opsets: 当自定义层定义一个symbolic的时候， 这边需要域也注册到相应的opts_version中； 
#比如g.op("ai.onnx.contrib::customlayer", intput, args_f=23.0, args_i=8)定义自定义层时，
# 则custom_opsets={"ai.onnx.contrib":11}, 11是前面opset_version设置保持一致
export_modules_as_functions=False
)

2. 自定义层处理

2.1 static 方式注册 - 通过Function方式的python方式进行自定义层定义

• 注意这个方式是将自己定义的函数，注册成一个onnx可以识别的operator

pytorch 版本 1.11
python=3.9
如果有m个输出,对应forward有多个返回值， 则设置为g.op(“ai.onnx.contrib::CustomLayer”, input, bias_i=8, scale_f=0.5, outputs=m)

示例代码

# 不带权重
class CustomLayer(Function):
    @staticmethod
    def forward(ctx, input, bias, scale):               
        out = input *scale + bias
        ctx.save_for_backward(torch.tensor(scale))
        return out]
        enable_onnx_checker
    @staticmethod
    def backward(ctx, gradout):
    	# write code
    	scale = ctx.saved_tensors
		return scale
		
    @staticmethod
    def