import torch
from transformers import BertModel, BertForMaskedLM
def convert_onnx():
model_path = ‘…/bertLM_model_pytorch/pytorch_model.pt’
dummy_input0 = torch.LongTensor(16, 128).to(torch.device(“cuda”))
dummy_input1 = torch.LongTensor(16, 128).to(torch.device(“cuda”))
dummy_input2 = torch.LongTensor(16, 128).to(torch.device(“cuda”))
dynamic_axes = {
‘dummy_input0’: {0: ‘batch_size’, 1: ‘seq_length’},
'dummy_input1 ': {0: ‘batch_size’, 1: ‘seq_length’},
‘dummy_input2’: {0: ‘batch_size’, 1: ‘seq_length’},
‘outputs’: {0: ‘batch_size’, 1: ‘seq_length’},
}
model = torch.load(model_path)
model.load_state_dict(loaded_model[‘state_dict’])
dummy_input =(dummy_input0, dummy_input1, dummy_input2)
onnx_path = ‘./pytorch_model.onnx’
torch.onnx.export(model,dummy_input, onnx_path,input_names=[‘dummy_input0’, 'dummy_input1 ',‘dummy_input2’], output_names=[‘outputs’],dynamic_axes=dynamic_axes)
print(‘convert retinaface to onnx finish!!!’)
I am trying to convert the pytorch model to onnx, and an error is reported after the above code runs:
AttributeError: ‘collections.OrderedDict’ object has no attribute ‘load_state_dict’
I want to know what is the cause of the error, thank you！