将 Pytorch 模型 .pth 转换为 onnx 模型

Question

我有一个预训练模型，格式为 .pth 扩展名。 我想把它转换成 Tensorflow protobuf。 但我没有找到任何方法来做到这一点。 我见过 onnx 可以将模型从 pytorch 转换为 onnx，然后从 onnx 转换为 Tensorflow。 但是使用这种方法，我在转换的第一阶段遇到了以下错误。

from torch.autograd import Variable
import torch.onnx
import torchvision
import torch 

dummy_input = Variable(torch.randn(1, 3, 256, 256))
model = torch.load('./my_model.pth')
torch.onnx.export(model, dummy_input, "moment-in-time.onnx")`

它给出了这样的错误。

File "t.py", line 9, in <module>
    torch.onnx.export(model, dummy_input, "moment-in-time.onnx")
  File "/usr/local/lib/python3.5/dist-packages/torch/onnx/__init__.py", line 75, in export
    _export(model, args, f, export_params, verbose, training)
  File "/usr/local/lib/python3.5/dist-packages/torch/onnx/__init__.py", line 108, in _export
    orig_state_dict_keys = model.state_dict().keys()
AttributeError: 'dict' object has no attribute 'state_dict'

可能的解决方案是什么？

Answer 1

尝试将您的代码更改为此

from torch.autograd import Variable

import torch.onnx
import torchvision
import torch

dummy_input = Variable(torch.randn(1, 3, 256, 256))
state_dict = torch.load('./my_model.pth')
model.load_state_dict(state_dict)
torch.onnx.export(model, dummy_input, "moment-in-time.onnx")

Answer 2

这意味着您的模型不是torch.nn.Modules类的子类。 如果将其作为子类，则应该可以使用。

Answer 3

这里的问题是您正在加载模型的权重，但是您也需要模型的架构，例如，如果您使用的是 mobilenet：

import torch
import torchvision.models as models

model=models.mobilenet_v3_large(weights)#Give your weights here
torch.onnx.export(model, torch.rand(1,3,640,640), "MobilenetV3.onnx")

有关更多信息，请参阅： https : //pytorch.org/tutorials/recipes/recipes/saving_and_loading_models_for_inference.html