Pytorch如何保存和加载模型参数

最新推荐文章于 2024-02-22 19:15:52 发布

爱不持久

最新推荐文章于 2024-02-22 19:15:52 发布

阅读量1w

点赞数 13

分类专栏： python相关学习卷积神经网络深度学习文章标签： python 深度学习网络

本文链接：https://blog.csdn.net/wacebb/article/details/108021921

版权

python相关学习同时被 3 个专栏收录

42 篇文章 4 订阅

订阅专栏

深度学习

32 篇文章 3 订阅

订阅专栏

卷积神经网络

26 篇文章 0 订阅

订阅专栏

pytorch 保存和加载模型的方法有两种：

1.保存网络的参数

import torch
#导入模块

net=Net()
#创建网络，当然还需要损失函数梯度等省略


PATH='state_dict_model.pth'
#先建立路径
torch.save(net.state_dict()，PATH)
#保存:可以是pth文件或者pt文件

model=Net()
model.load_state_dict(torch.load(PATH))
#载入保存的模型参数
model.eval()
#不启用 BatchNormalization 和 Dropout

2.保存整个网络

import torch

PATH = "entire_model.pt"
# Save
torch.save(net, PATH)

# Load
model = torch.load(PATH)
model.eval()

Remember too, that you must call model.eval() to set dropout and batch normalization layers to evaluation mode before running inference. Failing to do this will yield inconsistent inference results.