Pytorch从指定epoch恢复训练

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博客