[Pytorch] 模型的保存torch.save()与加载torch.load()

让我安静会

已于 2022-03-01 17:20:32 修改

阅读量2.9k

点赞数 4

分类专栏： model-pytorch 文章标签： pytorch 深度学习 python

于 2022-02-27 10:22:36 首次发布

本文链接：https://blog.csdn.net/qq_31225201/article/details/123160335

版权

51 篇文章 2 订阅

订阅专栏

在模型训练和验证之后，可以把训练好的模型进行保存。
后续进行测试的时候，可以直接加载模型进行测试。

将模型保存下来，如果后续有数据了可以继续训练，或者直接加载进行测试。

torch.save(model.state_dict(), PATH)  # 保存模型参数

torch.save(model,PATH)  # 保存整个模型

model = MyModel().to(device)
checkpoint = torch.load(config['save_path'])  # 先加载参数
model.load_state_dict(checkpoint)  # 再让模型加载参数, 恢复得到模型

model = torch.load(PATH)

在torch.load()中，常常会使用map_location进行cpu和gpu的转化。
（1）【GPU->CPU】
比如模型训练的时候是在GPU上进行并保存的，测试的时候却想在CPU上进行训练：

model = torch.load(PATH, map_location='cpu')

（2）【CPU->GPU】
转为GPU需要指明在哪块GPU上，例如转到第0块GPU：

model = torch.load(PATH, map_location=lambda storage, loc: storage.cuda(0))

（3）【GPU->GPU】
不同块GPU的转换，例如第1块转到第0块：

torch.load(PATH, map_location={'cuda:1':'cuda:0'})

Note：如果只保存了模型参数，就加载模型参数；如果保存了整个模型，就加载整个模型；上述两组一一对应。

参考：

pytorch 状态字典:state_dict使用详解：https://blog.csdn.net/Bruce_0712/article/details/111990905
torch.load_state_dict()函数的用法总结：https://blog.csdn.net/ChaoMartin/article/details/118686268
pytorch cpu与gpu load时相互转化 torch.load(map_location=)：https://blog.csdn.net/bc521bc/article/details/85623515

关注