如何保存与调用 pytorch 训练好的网络模型

问题描述

深度神经网络的模型训练起来十分困难，那么训练好的 pytorch 网络模型就应该把它保存起来，下次使用的时候直接调用就行了，应该怎么保存这个模型呢？

解决方法

Pytorch 主要提供了两种方法，分别是：保存模型参数的方法 和 保存整个模型的方法

方法一：只保存模型参数

#保存
torch.save(the_model.state_dict(), PATH)
#读取
the_model = TheModelClass(*args, **kwargs)
the_model.load_state_dict(torch.load(PATH))

方法二：保存整个模型

#保存
torch.save(the_model, PATH)
#读取
the_model = torch.load(PATH)

需要注意的是保存好的文件名推荐是 ×××.tar 格式的，例如PATH的格式可以设置为：'./model_file_name/the_model_name.tar'

 

还可以参看这篇文章：

https://blog.csdn.net/weixin_41680653/article/details/93768559