请选择 进入手机版 | 继续访问电脑版
查看: 577|回复: 0

pytorch保存模型等相关参数,利用torch.save(),以及读取保存...

[复制链接]

664

主题

1232

帖子

6445

积分

xdtech

Rank: 5Rank: 5

积分
6445
发表于 2020-5-13 08:36:37 | 显示全部楼层 |阅读模式
假设网络为model = Net(), optimizer = optim.Adam(model.parameters(), lr=args.lr), 假设在某个epoch,我们要保存模型参数,优化器参数以及epoch
一、
1. 先建立一个字典,保存三个参数:
state = {‘net':model.state_dict(), 'optimizer'ptimizer.state_dict(), 'epoch':epoch}
2.调用torch.save():
torch.save(state, dir)
其中dir表示保存文件的绝对路径+保存文件名,如'/home/qinying/Desktop/modelpara.pth'
二、
当你想恢复某一阶段的训练(或者进行测试)时,那么就可以读取之前保存的网络模型参数等。
checkpoint = torch.load(dir)
model.load_state_dict(checkpoint['net'])
optimizer.load_state_dict(checkpoint['optimizer'])
start_epoch = checkpoint['epoch'] + 1

回复

使用道具 举报

您需要登录后才可以回帖 登录 | 立即注册

本版积分规则

快速回复 返回顶部 返回列表