4.8 PyTorch模型的保存与加载

欢迎订阅本专栏:《PyTorch深度学习实践》
订阅地址:https://blog.csdn.net/sinat_33761963/category_9720080.html

  • 第二章:认识Tensor的类型、创建、存储、api等,打好Tensor的基础,是进行PyTorch深度学习实践的重中之重的基础。
  • 第三章:学习PyTorch如何读入各种外部数据
  • 第四章:利用PyTorch从头到尾创建、训练、评估一个模型,理解与熟悉PyTorch实现模型的每个步骤,用到的模块与方法。
  • 第五章:学习如何利用PyTorch提供的3种方法去创建各种模型结构。
  • 第六章:利用PyTorch实现简单与经典的模型全过程:简单二分类、手写字体识别、词向量的实现、自编码器实现。
  • 第七章:利用PyTorch实现复杂模型:翻译机(nlp领域)、生成对抗网络(GAN)、强化学习(RL)、风格迁移(cv领域)。
  • 第八章:PyTorch的其他高级用法:模型在不同框架之间的迁移、可视化、多个GPU并行计算。

方法一(推荐)

只保存和加载模型的参数

# 保存模型参数
def save_model(the_model, PATH):
    torch.save(the_model.state_dict(), PATH)
# 加载模型参数
def load_model(PATH):
    the_model = TheModelClass(*args, **kwargs)
    the_model.load_state_dict(torch.load(PATH))

方法二

在这种情况下,序列化的数据被绑定到特定的类和固定的目录结构,所以当在其他项目中使用时,或者在一些严重的重构器之后它可能会以各种方式break。

# 保存模型参数
def save_model(the_model, PATH):
    torch.save(the_model, PATH)
# 加载模型参数
def load_model(PATH):
    the_model = torch.load(PATH)
©️2020 CSDN 皮肤主题: 终极编程指南 设计师:CSDN官方博客 返回首页