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

可能的解决方案是什么？

- Urvish

---

你的.pth文件是一个状态字典，而不是完整的模型。首先需要创建一个模型，然后加载该状态字典，然后开始转换过程。请参考这个答案。 - layog

那种方法需要编写模型。但是我有预训练的模型，而且我不知道它的确切结构。因此，我不能像那个答案中那样定义模型。我该怎么办？ - Urvish

然后，确定架构就变得非常困难。你可以通过查看参数大小来猜测架构，但即使在查看大小后猜测正确的架构也非常困难，因为残差网络与非残差网络将具有相同大小的参数。你最好从预训练权重源获取架构定义。 - layog

好的。让我看看我能否做到。感谢您的帮助。如果我有.pth.tar文件，那么流程是否相同或有所改变？ - Urvish

Answer 1

这里的问题是您正在加载模型的权重，但您还需要在此处使用您的模型架构，例如如果您正在使用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

- Ashwinkumar J S

Answer 2

try changing your code to this

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

- faraz faraz

---

您的模型变量在哪里定义？ - Greg7000

这个答案不完整，模型变量未定义。 - cyrusbehr

Answer 3

这意味着你的模型不是torch.nn.Modules类的子类。如果你将它变成一个子类，就应该可以工作了。

- user129916

---

如何做到这一点？请提供一些指导。 - Urvish

当您定义类时，请编辑标题为：class ModelName(nn.Module)。 - user129916