在使用pytorch模型训练完成之后,我们现在使用的比较多的一种方法是将pytorch模型转成onnx格式的模型中间文件,然后再根据使用的硬件来生成具体硬件使用的深度学习模型,比如TensorRT。
在从pytorch模型转为onnx时,我们可能会遇到部分算子无法转换的问题,本篇注意记录下解决方法。
在导出onnx时,如果出现报错的算子,可以先在下面的链接中查找onnx算子是否支持
https://github.com/onnx/onnx/blob/main/docs/Operators.md
pytorch中有,onnx中也有的算子
导出时使用的onnx op 版本低导致
这个就好解决了,把op库的版本提高就行,但是有可能提高了版本以后,又出现了原来支持的算子现在又不支持了,这个再说
pytorch中没有注册某个onnx算子
如果是这种情况,就按照下面的方式进行:
from torch.onnx import register_custom_op_symbolic
# 创建一个asinh算子的symblic,符号函数,用来登记
# 符号函数内部调用g.op, 为onnx计算图添加Asinh算子
# g: 就是graph,计算图
# 也就是说,在计算图中添加onnx算子
# 由于我们已经知道Asinh在onnx是有实现的,所以我们只要在g.op调用这个op的名字就好了
# symblic的参数需要与Pytorch的asinh接口函数的参数对齐
# def asinh(input: Tensor, *, out: Optional[Tensor]=None) -> Tensor: ...
def asinh_symbolic(g, input, *, out=None):return g.op("Asinh", input)# 在这里,将asinh_symbolic这个符号函数,与PyTorch的asinh算子绑定。也就是所谓的“注册算子”
# asinh是在名为aten的一个c++命名空间下进行实现的# aten是"a Tensor Library"的缩写,是一个实现张量运算的C++库
register_custom_op_symbolic('aten::asinh', asinh_symbolic, 12)
另外一个写法
这个是类似于torch/onnx/symbolic_opset*.py中的写法
通过torch._internal中的registration来注册这个算子,让这个算子可以与底层C++实现的aten::asinh绑定
一般如果这么写的话,其实可以把这个算子直接加入到torch/onnx/symbolic_opset*.py中
import functools
from torch.onnx import register_custom_op_symbolic
from torch.onnx._internal import registration_onnx_symbolic = functools.partial(registration.onnx_symbolic, opset=9)@_onnx_symbolic('aten::asinh')
def asinh_symbolic(g, input, *, out=None):return g.op("Asinh", input)
pytorch中有,onnx中无的算子
继承torch.autograd.Function实现自定义算子
import torch
import torch.onnx
import onnxruntime
from torch.onnx import register_custom_op_symbolicOperatorExportTypes = torch._C._onnx.OperatorExportTypesclass CustomOp(torch.autograd.Function):@staticmethod def symbolic(g: torch.Graph, x: torch.Value) -> torch.Value:return g.op("custom_domain::customOp2", x)@staticmethoddef forward(ctx, x: torch.Tensor) -> torch.Tensor:ctx.save_for_backward(x)x = x.clamp(min=0)return x / (1 + torch.exp(-x))customOp = CustomOp.apply
然后再自己实现custom_domain::customOp2这个算子,如果用TensorRT,就需要自己实现一个插件。