参考https://blog.csdn.net/qq_37546267/article/details/106767640 利用下面代码将pth模型转换为onnx
import torch
from torch.autograd import Variable
import onnx
print(torch.__version__)
# torch --> onnx
input_name = ['input']
output_name = ['output']
input = Variable(torch.randn(1, 3, 224, 224)).cuda()
# model = torchvision.models.resnet50(pretrained=True).cuda()
model = torch.load('resnet50.pth', map_location="cuda:0")
torch.onnx.export(model, input, 'resnet50.onnx', input_names=input_name, output_names=output_name, verbose=True)
# 模型可视化
# netron.start('resnet50.onnx')
转换时出现 AttributeError: 'dict' object has no attribute 'training' 的问题。
通过查看https://pytorch.org/tutorials/advanced/super_resolution_with_onnxruntime.html 关于如何使用pytroch到处到onnx,和https://pytorch.org/tutorials/beginner/saving_loading_models.html 关于torch.save和torch.load的使用。找到了之前代码中使用的是
torch.save(model.state_dict(), PATH)
这种保存的是模型的字典。
model = torch.load(PATH)
再使用torch.load时,会报错。
所以,在保存模型是,使用:
torch.save(model, PATH)
再运行最上面的转换代码。可以正常工作。