將 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