【Pytorch】神经网络模型保存与加载

【Pytorch】神经网络模型保存与加载

神经网络模型的保存与加载有2种方法，一种方法会保存及加载整个网络（包括网络的图结构），另外一种方法只保存及加载网络的参数名及值（不包括网络的图结构）。

方法1

保存及加载整个网络，包括网络的图结构。代码如下：

# 保存网络
torch.save(model, PATH) # model是想要保存的模型，PATH是保存下来的文件的路径
# 加载网络
model2 = torch.load(PATH)

方法2

保存及加载网络的参数名及值，不包括网络的图结构。代码如下：

# 保存网络
torch.save(model.state_dict(), PATH) # model.state_dict()是想要保存的模型的参数名及值，PATH是保存下来的文件的路径
# 加载网络
model2 = MyModel() # 初始化准备载入参数的新模型
state_dict = torch.load(PATH) # 同样是torch.load(PATH)，在方法1中返回的是整个模型，在方法2中返回的是模型的参数及值。区别就在于当初PATH这个路径的文件里，保存的是整个模型还是只有模型的参数及值
model2.load_state_dict(state_dict)

参考：

源码详解Pytorch的state_dict和load_state_dict（只讲了1种方法，先save后load）

Pytorch保存和加载模型(load和load_state_dict)（讲了2种方法）