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.
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.
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.
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.