Pytorch无法将训练好的模型导出为ONNX。

Question

我一直在使用Pytorch框架训练一个模型，其中包含多个卷积层（3x3、步长1、padding same）。这个模型表现良好，我想在Matlab中将其用于推断。为此，NN框架之间交换的ONNX格式似乎是（唯一？）解决方案。可以使用以下命令导出模型：

torch.onnx.export(net.to('cpu'), test_input,'onnxfile.onnx')

这是我的卷积神经网络结构定义：

class Encoder_decoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = nn.Sequential(
        nn.Conv2d(2,8, (3, 3),stride = 1, padding='same'),
        nn.ReLU(),
        nn.Conv2d(8,8, (3, 3),stride = 1, padding='same'),
        nn.ReLU(),
        nn.Conv2d(8,16, (3, 3),stride = 1, padding='same'),
        nn.ReLU(),
        nn.Conv2d(16,16, (3, 3),stride = 1, padding='same'),
        nn.ReLU(),
        nn.Conv2d(16,32, (3, 3),stride = 1, padding='same'),
        nn.ReLU(),
        nn.Conv2d(32,32, (3, 3),stride = 1, padding='same'),
        nn.ReLU(),
        nn.Conv2d(32,64, (3, 3),stride = 1, padding='same'),
        nn.ReLU(),
        nn.Conv2d(64,64, (3, 3),stride = 1, padding='same'),
        nn.ReLU(),
        nn.Conv2d(64,128, (3, 3),stride = 1, padding='same'),
        nn.ReLU(),
        nn.Conv2d(128,128, (3, 3),stride = 1, padding='same'),
        nn.ReLU(),
        nn.Conv2d(128,1, (1, 1))
        )


    def forward(self, x):
        x = self.model(x)
        
        return x

然而，当我运行torch.onnx.export命令时，我得到以下错误：

RuntimeError: Exporting the operator _convolution_mode to ONNX opset version 9 is not supported. Please feel free to request support or submit a pull request on PyTorch GitHub.

我尝试更改opset，但没有解决问题。ONNX完全支持卷积神经网络。此外，我正在google colab中训练网络。

您知道将模型转移到matlab的其他方法吗？

- Mahindra Many

Answer 1

目前，PyTorch中不支持_convolution_mode操作符（点击查看）。这是由于使用了padding='same'。

您需要将padding更改为整数值或更改为其等效值。请参考Pytorch中的Same padding equivalent。

- kHarshit

Answer 2

我想出了一个变通方法：

...
def calc_same_padding(kernel_size, stride, input_size):
    if isinstance(kernel_size, Sequence):
        kernel_size = kernel_size[0]

    if isinstance(stride, Sequence):
        stride = stride[0]

    if isinstance(input_size, Sequence):
        input_size = input_size[0]

    pad = ((stride - 1) * input_size - stride + kernel_size) / 2
    return int(pad)

def replace_conv2d_with_same_padding(m: nn.Module, input_size=512):
    if isinstance(m, nn.Conv2d):
        if m.padding == "same":
            m.padding = calc_same_padding(
                kernel_size=m.kernel_size,
                stride=m.stride,
                input_size=input_size
            )

...
model = MyModel()

model.apply(lambda m: replace_conv2d_with_same_padding(m, 512))
example_input = torch.ones((1, 3, 512, 512))

torch.onnx.export(model,
                  example_input,
                  input_names=["input"],
                  output_names=["output"],
                  f=save_path,
                  opset_version=12)

所有我的输入/输出张量都有偶数维度，例如512x512/256x256/128x128等，因此这里的输入尺寸并不重要。

- RedEyed