Transferring pretrained pytorch model to onnx

Question

I am trying to convert pytorch model to ONNX, in order to use it later for TensorRT. I followed the following tutorial https://pytorch.org/tutorials/advanced/super_resolution_with_caffe2.html, but my kernel dies all the time.

This is the code that I implemented.

 # Some standard imports
import io
import numpy as np

from torch import nn
import torch.onnx
from deepformer.nets.quicknat import quickNAT
param = {
    'num_channels': 64,
    'num_filters': 64,
    'kernel_h': 5,
    'kernel_w': 5,
    'kernel_c': 1,
    'stride_conv': 1,
    'pool': 2,
    'stride_pool': 2,
    'num_classes': 1,
    'padding': 'reflection'
}


net = quickNAT(param)
checkpoint_path = 'checkpoint_epoch36_loss0.78.t7'
checkpoints=torch.load(checkpoint_path)
map_location = lambda storage, loc: storage
if torch.cuda.is_available():
    map_location = None
net.load_state_dict(checkpoints['net'])
net.train(False)
# Input to the modelvcdfx  
x = torch.rand(1, 64, 256, 1600, requires_grad=True)

# Export the model
torch_out = torch.onnx._export(net,             # model being run
                               x,                       # model input (or a tuple for multiple inputs)
                               "quicknat.onnx", # where to save the model (can be a file or file-like object)
                               export_params=True)      # store the trained parameter weights inside the model file

Answer 1

What is the output you get? It seems SuperResolution is supported with the export operators in pytorch as mentioned in the documentation

Are you sure the input to your model is:

x = torch.rand(1, 64, 256, 1600, requires_grad=True)

That could be the variable that you used for training, since for deployment you run the network on one or multiple images the dummy input to export to onnx is usually:

dummy_input = torch.randn(1, 3, 720, 1280, device='cuda')

With 1 being the batch size, 3 being the channels of the image(RGB), and then the size of the image, in this case 720x1280. Check on that input, I guess you don't have a 64 channel image as input right?

Also, it'd be helpful if you post the terminal output to see where it fails. Good luck!