使用Pytorch导出自定义ONNX算子的示例代码
作者：太阳花的小绿豆
在实际部署模型时有时可能会遇到想用的算子无法导出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 OperatorExportTypes
class CustomAffineGrid(Function):
@staticmethod
def forward(ctx, theta: torch.Tensor, size: torch.Tensor):
grid = F.affine_grid(theta=theta, size=size.cpu().tolist())
return grid
@staticmethod
def 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, grid=grid, mode="bilinear", padding_mode="zeros")
return x
def 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(model=custum_model,
args=(x, theta, size),
f="custom.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_version=16,
operator_export_type=OperatorExportTypes.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 OperatorExportTypes
class CustomRot90AndScale(Function):
@staticmethod
def forward(ctx, x: torch.Tensor):
x = torch.rot90(x, k=1, dims=(3, 2)) # clockwise 90
x *= 1.2
return x
@staticmethod
def symbolic(g: torch.Graph, x: torch.Tensor):
return g.op("Rot90AndScale", x, k_i=1, scale_f=1.2, clockwise_s="yes")
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(model=custum_model,
args=(x,),
f="custom_rot90.onnx",
input_names=["input"],
output_names=["output"],
dynamic_axes={"input": {2: "h0", 3: "w0"},
"output": {2: "w0", 3: "h0"}},
opset_version=16,
operator_export_type=OperatorExportTypes.ONNX_FALLTHROUGH)
if __name__ == '__main__':
main()使用netron打开自己导出的onnx文件，可以看到如下所示信息。
到此这篇关于使用Pytorch导出自定义ONNX算子的文章就介绍到这了,更多相关使用Pytorch导出自定义ONNX算子内容请搜索脚本之家以前的文章或继续浏览下面的相关文章希望大家以后多多支持脚本之家！