Export Pytorch’s nn.MaxUnpool2d to ONNX

This artcile explains how to export Pytorch’s nn.MaxUnpool2d to ONNX.

About Pytorch’s MaxUnpool2d and ONNX’s Unpool2d

MaxUnpool2d is the inverse operation of MaxPool2d, it can be used to increase the resolution of a feature map. The corresponding operator in ONNX is Unpool2d, but it cannot be simply exported from Pytorch because the indices specification is different for both operators.

Therefore, if you try to export a model containing MaxUnpool2d to ONNX, you will get the following error.

RuntimeError: Exporting the operator max_unpool2d to ONNX opset version 11 is not supported. Please open a bug to request ONNX export support for the missing operator.

How to fix the MaxUnpool2d export problem

To be able to export MaxUnpool2d, replace Pytorch’s nn.MaxUnpool2d with the implementation from mmediting. The mmediting implementation defines a symbolic method which combines ONNX operators to compute Unpool2d indices.

open-mmlab/mmediting

Use the above code and replace it with as follows.

#self.unpool = nn.MaxUnpool2d(2, 2)  #torch
self.unpool = MaxUnpool2d(2, 2) #onnx

This way the export will be successful.

Confirmation of the export result

Let’s export Deep-Image-Matting-Pytorch which contains the MaxUnpool2d operator.

foamliu/Deep-Image-Matting-PyTorch

Apply the above modifications to model.py and export using test.py

Then make sure both the original Pytorch model and the exported ONNX model give the same inference results.

nn.MaxUnpool2d

ONNX (mmediting.MaxUnpool2d)

ax Inc. has developed ailia SDK, which enables cross-platform, GPU-based rapid inference.

ax Inc. provides a wide range of services from consulting and model creation, to the development of AI-based applications and SDKs. Feel free to contact us for any inquiry.