pytorch中模型的保存与加载

最新推荐文章于 2024-06-02 23:52:21 发布

sugarzwp

最新推荐文章于 2024-06-02 23:52:21 发布

阅读量1k

点赞数 1

文章标签： pytorch

本文链接：https://blog.csdn.net/qq_41883764/article/details/90703597

版权

特别开心的训练出来10个模型，准备加载画confusion_matrix…然后就遇到了下面的问题：
AttributeError: ‘VGG’ object has no attribute 'copy’
原来是由于下面这个原因：模型的保存与加载的方式不对应！

1.torch.save()实现对网络结构和模型参数的保存。有两种保存方式：一是保存整个神经网络的的结构信息和模型参数信息，save的对象是网络net；二是只保存神经网络的训练模型参数，save的对象是net.state_dict()。

    torch.save(net1, '7-net.pth')                     # 保存整个神经网络的结构和模型参数  
    torch.save(net1.state_dict(), '7-net_params.pth') # 只保存神经网络的模型参数

我的代码中是这样写的：
'net': net.state_dict() if use_cuda else net
如果用gpu的话，就只保存网络的训练模型参数
如果用cpu的话，就保存整个网络的结构信息和模型参数信息

2.对应上面两种保存方式，重载方式也有两种。对应第一种完整网络结构信息，重载的时候通过checkpoint=torch.load(‘.pth’),net=checkpoint[‘net’]直接初始化新的神经网络对象即可。对应第二种只保

关注