I am trying to convert PyTorch model to Onnx model.
I keep ending up with the following error:
Traceback (most recent call last):
File "/tmp/pycharm_project_896/utils/pytorch2onnx.py", line 15, in <module>
for k, v in state_dict.items():
File "/usr/local/lib/python2.7/dist-packages/torch/nn/modules/module.py", line 398, in __getattr__
type(self).__name__, name))
AttributeError: 'DataParallel' object has no attribute 'items'
I have been using the following code:
import torch
from torch.autograd import Variable
# Load the trained model from file
path = "model.pth"
#pytorch_network = torch.load(path)
state_dict = torch.load("model.pth")
#pytorch_network.state_dict()
# Create new OrderedDict that does not contain the module
from collections import OrderedDict
new_state_dict = OrderedDict()
for k, v in state_dict.items():
name = k[7:] # remove module
new_state_dict[name] = v
pytorch_network.load_state_dict(new_state_dict)
# Export the trained model to ONNX
#pytorch_network = torch.nn.DataParallel(pytorch_network, device_ids=range(torch.cuda.device_count()))
dummy_input = Variable(torch.randn(1, 3, 224, 224))
torch.onnx.export(pytorch_network, dummy_input, "model.onnx")
Can anyone tell me where I might be going wrong. I keep ended up with the same error.
I have referred this Issue too but I don’t seem to end with a solution.
Thanks