飞凡 做电商网站,沈阳市和平区网站建设,网站建设注册,湖南省郴州市汝城县在实际部署模型时有时可能会遇到想用的算子无法导出onnx#xff0c;但实际部署的框架是支持该算子的。此时可以通过自定义onnx算子的方式导出onnx模型#xff08;注#xff1a;自定义onnx算子导出onnx模型后是无法使用onnxruntime推理的#xff09;。下面给出个具体应用中的…在实际部署模型时有时可能会遇到想用的算子无法导出onnx但实际部署的框架是支持该算子的。此时可以通过自定义onnx算子的方式导出onnx模型注自定义onnx算子导出onnx模型后是无法使用onnxruntime推理的。下面给出个具体应用中的示例需要导出pytorch的affine_grid算子但在pytorch的2.0.1版本中又无法正常导出该算子故可通过如下自定义算子代码导出。
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Function
from torch.onnx import OperatorExportTypesclass CustomAffineGrid(Function):staticmethoddef forward(ctx, theta: torch.Tensor, size: torch.Tensor):grid F.affine_grid(thetatheta, sizesize.cpu().tolist())return gridstaticmethoddef symbolic(g: torch.Graph, theta: torch.Tensor, size: torch.Tensor):return g.op(AffineGrid, theta, size)class MyModel(nn.Module):def __init__(self) - None:super().__init__()def forward(self, x: torch.Tensor, theta: torch.Tensor, size: torch.Tensor):grid CustomAffineGrid.apply(theta, size)x F.grid_sample(x, gridgrid, modebilinear, padding_modezeros)return xdef main():with torch.inference_mode():custum_model MyModel()x torch.randn(1, 3, 224, 224)theta torch.randn(1, 2, 3)size torch.as_tensor([1, 3, 512, 512])torch.onnx.export(modelcustum_model,args(x, theta, size),fcustom.onnx,input_names[input0_x, input1_theta, input2_size],output_names[output],dynamic_axes{input0_x: {2: h0, 3: w0},output: {2: h1, 3: w1}},opset_version16,operator_export_typeOperatorExportTypes.ONNX_FALLTHROUGH)if __name__ __main__:main()
在上面代码中通过继承torch.autograd.Function父类的方式实现导出自定义算子继承该父类后需要用户自己实现forward以及symbolic两个静态方法其中forward方法是在pytorch正常推理时调用的函数而symbolic方法是在导出onnx时调用的函数。对于forward方法需要按照正常的pytorch语法来实现其中第一个参数必须是ctx但对于当前导出onnx场景可以不用管它后面的参数是实际自己传入的参数。对于symbolic方法的第一个必须是g后面的参数任为实际自己传入的参数然后通过g.op方法指定具体导出自定义算子的名称以及输入的参数注上面示例中传入的都是Tensor所以可以直接传入对与非Tensor的参数可见下面一个示例。最后在使用时直接调用自己实现类的apply方法即可。使用netron打开自己导出的onnx文件可以看到如下所示网络结构。
有时按照使用的推理框架导出自定义算子时还需要设置一些参数非Tensor那么可以参考如下示例例如要导出int型的参数k那么可以通过传入k_i来指定要导出float型的参数scale那么可以通过传入scale_f来指定要导出string型的参数clockwise那么可以通过传入clockwise_s来指定
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Function
from torch.onnx import OperatorExportTypesclass CustomRot90AndScale(Function):staticmethoddef forward(ctx, x: torch.Tensor):x torch.rot90(x, k1, dims(3, 2)) # clockwise 90x * 1.2return xstaticmethoddef symbolic(g: torch.Graph, x: torch.Tensor):return g.op(Rot90AndScale, x, k_i1, scale_f1.2, clockwise_syes)class MyModel(nn.Module):def __init__(self) - None:super().__init__()def forward(self, x: torch.Tensor):return CustomRot90AndScale.apply(x)def main():with torch.inference_mode():custum_model MyModel()x torch.randn(1, 3, 224, 224)torch.onnx.export(modelcustum_model,args(x,),fcustom_rot90.onnx,input_names[input],output_names[output],dynamic_axes{input: {2: h0, 3: w0},output: {2: w0, 3: h0}},opset_version16,operator_export_typeOperatorExportTypes.ONNX_FALLTHROUGH)if __name__ __main__:main()
使用netron打开自己导出的onnx文件可以看到如下所示信息。
