如何获取未知PyTorch模型的输入张量形状 [英] How to get input tensor shape of an unknown PyTorch model

问题描述

我正在编写一个python脚本，该脚本将任何流行的框架（TensorFlow，Keras，PyTorch）的深度学习模型转换为ONNX格式。目前，我已将 tf2onnx 用于tensorflow和 keras2onnx 可以将keras转换为ONNX，并且可以使用。

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.

现在PyTorch集成了ONNX支持，因此我可以直接从PyTorch保存ONNX模型。但是问题是我需要为该模型输入张量形状，以便将其保存为ONNX格式。您可能已经猜到了，我正在编写此脚本来转换未知的深度学习模型。

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.

> 此处 是PyTorch的ONNX转换教程。上面写着：

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

限制¶
ONNX导出器是基于跟踪的导出器，这意味着它可以运行通过执行一次模型并导出在此运行期间实际运行的运算符。这意味着如果您的模型是动态的，例如根据输入数据更改行为，则导出将不准确。

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, e.g., 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

那么我怎么知道那个未知的PyTorch模型的输入张量的INPUT_SHAPE？还是有其他方法可以将PyTorch模型转换为ONNX？

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?

推荐答案

您可以以此为起点进行调试

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

之后，得到N，C并通过专门放置H来创建张量，w像这个玩具示例一样，没有像

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')

这篇关于如何获取未知PyTorch模型的输入张量形状的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！