PyTorch保存模型断点以及加载断点继续训练

在训练神经网络时，用到的数据量可能很大，训练周期较长，如果半途中断了训练，下次从头训练就会很费时间，这时我们就想断点续训。

一、神经网络模型的保存，基本两种方式：
1. 保存完整模型model，torch.save(model, save_path)

2. 只保存模型的参数，torch.save(model.state_dict(), save_path)，多卡训练的话，在保存参数时，使用 model.module.state_dict( ) 。

二、保存模型的断点checkpoint

断点dictionary中一般保存训练的网络的权重参数、优化器的状态、学习率lr_scheduler 的状态以及epoch 。

checkpoint = {'parameter': model.module.state_dict(),              'optimizer': optimizer.state_dict(),              'scheduler': scheduler.state_dict(),              'epoch': epoch}

torch.save(checkpoint, ‘./models/checkpoint/ckpt_{}.pth’.format(epoch+1))

三、加载断点继续训练

if resume: # Trueload_ckpt = torch.load(ckpt_dir, map_location=device)                                 # 从断点路径加载断点，指定加载到CPU内存或GPUload_weights_dict = {k: v for k, v in load_ckpt['parameter'].items()                                      if model.state_dict()[k].numel() == v.numel()}  # 简单验证model.load_state_dict(load_weights_dict, strict=False) 
# 如果是多卡训练，加载weights后要设置DDP模式，然后先定义一下optimizer和scheduler，之后再加载断点中保存的optimizer和scheduler以及设置epoch，
optimizer.load_state_dict(load_ckpt['optimizer'])                                     # 加载优化器状态scheduler.load_state_dict(load_ckpt['scheduler'])                                     # 加载scheduler状态
start_epoch = load_ckpt['epoch']+1                                                    # 设定继续训练的epoch起点iter_epochs = range(start_epoch, args.epochs)                                         # arg.epochs指出训练的总epoch数，包括断点前的训练次数

Enjoy it!