How To Convert Deform Conv2d into onnx format in PyTorch?

Yanwei Liu

·

Follow

1 min read

·

Jan 2, 2024

--

Listen

Share

Installation

pip install deform_conv2d_onnx_exporter

Usage

import torch
import torch.onnx 
import pwd
import os
from arch_network import DeformConvNetwork
from torchvision.ops.deform_conv import DeformConv2d
import deform_conv2d_onnx_exporter
deform_conv2d_onnx_exporter.register_deform_conv2d_onnx_op()

"""
Convert model that contains deform conv2d layer into onnx format
https://pypi.org/project/deform-conv2d-onnx-exporter/1.0.0/
"""


#Function to Convert to ONNX 
def Convert_ONNX(model): 

    # set the model to inference mode 
    model.eval() 

    # Let's create a dummy input tensor  
    dummy_input = torch.randn(1, 3, 256, 256, requires_grad=True)  

    # Export the model   
    torch.onnx.export(model,         # model being run 
         dummy_input,       # model input (or a tuple for multiple inputs) 
         "./onnx/network.onnx",       # where to save the model  
         export_params=True,  # store the trained parameter weights inside the model file 
         opset_version=12,    # the ONNX version to export the model to 
         do_constant_folding=True,  # whether to execute constant folding for optimization 
         input_names = ["input", "offset"],   # the model's input names 
         output_names = ['output'], # the model's output names 
         dynamic_axes={'modelInput' : {0 : 'batch_size'},    # variable length axes 
                                'modelOutput' : {0 : 'batch_size'}}) 
    print(" ") 
    print('Model has been converted to ONNX') 
if __name__ == "__main__": 
    model = DeformConvNetwork()
    model.load_state_dict(torch.load(f"/home/{pwd.getpwuid(os.getuid()).pw_name}/experiments/network.pth"))
    Convert_ONNX(model)