使用torch.load_state_dict()将pth转onnx模型出现的问题

model = ResNet34(num_classes=20)
model.load_state_dict(torch.load(model_path, map_location=device),strict=False)
model.to(device=device)
model.eval() #设置模型为推理模式(重要

错误提示:
torch.nn.modules.module.ModuleAttributeError: ‘Resnet34‘ object has no attribute ‘copy‘

发现问题出现在训练模型的代码中,在保存模型时,使用的是

torch.save(model, path)

保存模型结构 和 参数
故加载的时候应该

model = torch.load(model_path, map_location=device)
model.to(device=device)
model.eval()

===============================================================================
最近加载函数的时候发生了错误,提示:torch.nn.modules.module.ModuleAttributeError: 'Unet' object has no attribute 'copy'

原因可能是对参数保存和读取没有掌握,保存和加载有两种方式,一种是全部打包,另一种是只保存参数,

*保存和加载整个模型 *

torch.save(model_object, 'model.pth')  
model = torch.load('model.pth')  


# 仅保存和加载模型参数  
torch.save(model_object.state_dict(), 'params.pth')  
model_object.load_state_dict(torch.load('params.pth'))  
个人建议两个都分别保存一下,需要的时候都能用上

torch.save(model, 'net.pth')

torch.save(model.state_dict(), "param_net.pth")

我的原因就是保存了全部,加载的时候用的方式却是只加载参数的方式,所以应该改为

net = torch.load(model_weights, map_location=device)
net.to(device=device)

————————————————
版权声明:本文为CSDN博主「MurasameFan」的原创文章,遵循CC 4.0 BY-SA版权协议,转载请附上原文出处链接及本声明。
原文链接:https://blog.csdn.net/qq_34419607/article/details/110347137