pytorch转onnx 手动添加不支持的算子-摩杜云开发者社区

PyTorch转ONNX：手动添加不支持的算子

介绍

PyTorch是一个流行的开源深度学习框架，它提供了丰富的预训练模型和灵活的计算图构建机制。然而，有时候我们需要将PyTorch模型转换为ONNX格式，以便在其他平台上运行。ONNX是一种开放的深度学习框架交换格式，它允许我们在不同的框架之间共享和部署模型。

在PyTorch中，我们可以使用torch.onnx.export()函数将模型转换为ONNX格式。然而，有些PyTorch算子可能没有与之对应的ONNX算子，这就需要我们手动添加不支持的算子。本文将介绍如何手动添加不支持的算子，以确保成功转换PyTorch模型为ONNX格式。

转换流程

导入所需的库

import torch
import torch.onnx
import onnx
from onnx import helper, shape_inference

定义PyTorch模型

class MyModel(torch.nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.linear = torch.nn.Linear(2, 1)

    def forward(self, x):
        return self.linear(x)

创建一个示例输入

dummy_input = torch.randn(1, 2)

导出PyTorch模型为ONNX格式

torch.onnx.export(
    MyModel(),
    dummy_input,
    "model.onnx",
    verbose=True
)

在这一步中，如果遇到不支持的算子，会抛出异常。

手动添加不支持的算子

如果在转换过程中遇到不支持的算子，我们需要手动添加它们到ONNX模型中。首先，我们可以使用torch.onnx.export()函数的operator_export_type参数将模型导出为一个具有中间表示的ONNX文件，而不是最终的ONNX文件。

torch.onnx.export(
    MyModel(),
    dummy_input,
    "model.onnx",
    verbose=True,
    operator_export_type=torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK
)

导出的中间表示ONNX文件中，所有不支持的算子将以ATen算子的形式存在。我们可以使用ONNX库的shape_inference.infer_shapes()函数推断这些ATen算子的输出形状。

model = onnx.load("model.onnx")
model = shape_inference.infer_shapes(model)

现在，我们可以手动添加不支持的算子到模型中。首先，我们需要了解不支持的算子的输入和输出形状。然后，我们可以使用ONNX库的helper.make_node()函数创建一个新的算子节点。

node = helper.make_node(
    'CustomOp',
    inputs=['input'],
    outputs=['output'],
    attribute=dict(attribute_name='value')
)

最后，我们需要将新创建的算子节点添加到模型的图中，并保存为最终的ONNX文件。

model.graph.node.extend([node])
onnx.save(model, "custom_model.onnx")

注意，CustomOp需要根据实际情况替换为具体的自定义算子。

结论

通过手动添加不支持的算子，我们可以成功转换PyTorch模型为ONNX格式。这使得我们能够在其他平台上部署和运行模型，从而获得更好的灵活性和可移植性。然而，我们需要确保手动添加的算子与原始算子具有相同的功能和正确的输入输出形状。