weixin_44725059的博客

原创 pytorch的模型保存和加载方式汇总

torch.save函数有两种保存方式：一种是保存整个模型，此时模型的type应该为继承自nn.Module的类，这里则为类LeNet5；另一种是仅保存模型的参数，此时模型的type应该为有序字典即类OrderedDict。torch.save函数将序列化的对象保存到磁盘。通过pickle可以保存各种对象的模型、张量和字典。保存的模型一般是以.pt/.pth/.pkl文件后缀，作用都一样。将为你以后恢复模型提供最大的灵活性，推荐方法。使用torch.save函数保存模型的。state_dict保存。