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.