使用Pytorch導出自定義ONNX算子的示例代碼

更新時間：2024年03月08日 12:00:30 作者：太陽花的小綠豆

這篇文章主要介紹了使用Pytorch導出自定義ONNX算子的示例代碼,下面給出個具體應用中的示例：需要導出pytorch的affine_grid算子,但在pytorch的2.0.1版本中又無法正常導出該算子,故可通過如下自定義算子代碼導出,需要的朋友可以參考下

在實際部署模型時有時可能會遇到想用的算子無法導出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父類的方式實現(xiàn)導出自定義算子，繼承該父類后需要用戶自己實現(xiàn)forward以及symbolic兩個靜態(tài)方法，其中forward方法是在pytorch正常推理時調(diào)用的函數(shù)，而symbolic方法是在導出onnx時調(diào)用的函數(shù)。對于forward方法需要按照正常的pytorch語法來實現(xiàn)，其中第一個參數(shù)必須是ctx但對于當前導出onnx場景可以不用管它，后面的參數(shù)是實際自己傳入的參數(shù)。對于symbolic方法的第一個必須是g，后面的參數(shù)任為實際自己傳入的參數(shù)，然后通過g.op方法指定具體導出自定義算子的名稱，以及輸入的參數(shù)（注：上面示例中傳入的都是Tensor所以可以直接傳入，對與非Tensor的參數(shù)可見下面一個示例）。最后在使用時直接調(diào)用自己實現(xiàn)類的apply方法即可。使用netron打開自己導出的onnx文件，可以看到如下所示網(wǎng)絡結構。

有時按照使用的推理框架導出自定義算子時還需要設置一些參數(shù)（非Tensor）那么可以參考如下示例，例如要導出int型的參數(shù)k那么可以通過傳入k_i來指定，要導出float型的參數(shù)scale那么可以通過傳入scale_f來指定，要導出string型的參數(shù)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文件，可以看到如下所示信息。