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

一、如何保存模型参数,优化器参数

 假设在某个epoch,我们要保存模型参数,优化器参数以及epoch。

1. 先建立一个字典,保存三个参数:

if steps % 10000 == 0:
            file_path = os.path.join(args.weight_dir, 'weights-{:07d}.pth.tar'.format(steps))
            state = {
            'epoch' :epoch+1,
            'state_dict' :model.state_dict(),
            'optimizer' : optimizer.state_dict(),
            }
            torch.save(state, file_path)

其中dir表示保存文件的绝对路径+保存文件名,如'/home/qinying/Desktop/modelpara.pth'

2、

当你想恢复某一阶段的训练(或者进行测试)时,那么就可以读取之前保存的网络模型参数等。

checkpoint = torch.load(dir)

model.load_state_dict(checkpoint['net'])

optimizer.load_state_dict(checkpoint['optimizer'])

start_epoch = checkpoint['epoch'] + 1

猜你喜欢

转载自blog.csdn.net/Scwabc_123/article/details/84866471
今日推荐