1、在训练时保存每个epoch的模型&optimizer&epoch
checkpoint= {'model': model.state_dict(), 'optimizer': optimizer.state_dict(), 'epoch': epoch}
torch.save(checkpoint, path)
2、从指定epoch恢复
path_checkpoint = "./models/checkpoint/ckpt_best_1.pth" # 断点路径
checkpoint = torch.load(path_checkpoint) # 加载断点
model.load_state_dict(checkpoint['net']) # 加载模型可学习参数
optimizer.load_state_dict(checkpoint['optimizer']) # 加载优化器参数
start_epoch = checkpoint['epoch'] # 设置开始的epoch
注意:lr_scheduler是根据内置的index进行更新的,外界的epoch对其并不会产生影响,因此要自己设置lr_scheduler的last_epoch;如果恢复时optimizer的param_groups(list of dicts)没有initial_lr,也要自己设置initial_lr(可以在训练时也保存下来每个epoch的lr),否则会报错,例如
optimizer = optim.AdamW([{'params': model.parameters(), 'initial_lr': config.lr}], lr=config.lr)
scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=[50, 120], gamma=0.5, last_epoch=config.iter)
参考:
pytorch 从指定epoch恢复训练 - 简书 (jianshu.com)
坑中坑之Pytorch断点重新训练/恢复训练scheduler设置/last_epoch研究_pytorch恢复训练_Mr.July的博客-CSDN博客