Преобразование модели 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. Если вы сделаете это подклассом, это должно работать.