How to get input tensor shape of an unknown PyTorch model

Question

I am writing a python script, which converts any deep learning models from popular frameworks (TensorFlow, Keras, PyTorch) to ONNX format. Currently I have used tf2onnx for tensorflow and keras2onnx for keras to ONNX conversion, and those work.

Now PyTorch has integrated ONNX support, so I can save ONNX models from PyTorch directly. But the problem is I will need input tensor shape for that model, in order to save it in ONNX format. As you already might have guessed, I am writing this script to convert unknown deep learning models.

Here is PyTorch's tutorial for ONNX conversion. There it's written:

Limitations ¶ The ONNX exporter is a trace-based exporter, which means that it operates by executing your model once, and exporting the operators which were actually run during this run. This means that if your model is dynamic, eg, changes behavior depending on input data, the export won't be accurate.

Similarly, a trace is might be valid only for a specific input size (which is one reason why we require explicit inputs on tracing). Most of the operators export size-agnostic versions and should work on different batch sizes or input sizes. We recommend examining the model trace and making sure the traced operators look reasonable.

The code snippet I am using is this:

import torch

def convert_pytorch2onnx(self):
    """pytorch -> onnx"""

    model = torch.load(self._model_file_path)

    # Don't know how to get this INPUT_SHAPE
    dummy_input = torch.randn(INPUT_SHAPE)
    torch.onnx.export(model, dummy_input, self._onnx_file_path)
    return

So how do I know the INPUT_SHAPE of the input tensor of that unknown PyTorch model? Or is there any other way to convert the PyTorch model to ONNX?

Answer 1

you can follow this as a starting point to debug

list(model.parameters())[0].shape # weights of the first layer in the format (N,C,Kernel dimensions) # 64, 3, 7 ,7

after that get the N,C and create a tensor out of that by specially putting H,W as None like this toy example

import torch
import torchvision

net = torchvision.models.resnet18(pretrained = True)

shape_of_first_layer = list(net.parameters())[0].shape #shape_of_first_layer

N,C = shape_of_first_layer[:2]

dummy_input = torch.Tensor(N,C)

dummy_input = dummy_input[...,:, None,None] #adding the None for height and weight

torch.onnx.export(net, dummy_input, './alpha')