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

最新推荐文章于 2025-09-06 23:40:27 发布

原创最新推荐文章于 2025-09-06 23:40:27 发布 · 996 阅读

6 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #神经网络 #深度学习

文章介绍了在Pytorch中如何保存和加载神经网络模型，包括保存整个网络结构和仅保存参数两种方式。方法一使用`torch.save(model,PATH)`保存和`torch.load(PATH)`加载整个模型；方法二利用`model.state_dict()`保存参数，然后通过`model2.load_state_dict(state_dict)`加载到新初始化的模型中。

部署运行你感兴趣的模型镜像

【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)