pytorch保存模型的坑

chenf99

于 2020-02-15 21:28:40 发布

阅读量5k

点赞数 6

CC 4.0 BY-SA版权

分类专栏：深度学习

本文链接：https://blog.youkuaiyun.com/chenf1999/article/details/104333709

本文介绍了PyTorch在保存和加载模型过程中遇到的设备错误和模型错误，包括如何处理模型默认在GPU上的问题，解决不同PyTorch版本之间的兼容性问题，以及如何正确保存和恢复优化器的状态。总结建议在保存时将模型和优化器参数转到CPU，并根据训练需求决定是否保存优化器状态。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

pytorch中保存模型相关的函数有3个：

有两种方法保存模型：

1. torch.save(model, path) # 直接保存整个模型
2. torch.save(model.state_dict(), path) # 保存模型的参数

相应地有两种方法加载保存的模型：

1. model = torch.load(path) # 直接加载模型
2. model = Model()                         # 先初始化一个模型
   model.load_state_dict(torch.load(path))

200万优质内容无限畅学