如何获取未知 PyTorch 模型的输入张量形状

Question

我正在编写一个 python 脚本，它将任何深度学习模型从流行的框架（TensorFlow、Keras、PyTorch）转换为 ONNX 格式。目前我已经使用 tf2onnx for tensorflow and keras2onnx 将 keras 转换为 ONNX，并且这些工作正常。

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

Here是PyTorch的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.

我使用的代码片段是这样的：

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？

Answer 1

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

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 作为 None 像这个玩具示例一样从中创建一个张量

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 模型的输入张量形状

How to get input tensor shape of an unknown PyTorch model

python

machine-learning

deep-learning

pytorch

onnx