在训练神经网络时,用到的数据量可能很大,训练周期较长,如果半途中断了训练,下次从头训练就会很费时间,这时我们就想断点续训。
一、神经网络模型的保存,基本两种方式:
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: # True load_ckpt = torch.load(ckpt_dir, map_location=device) load_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)
optimizer.load_state_dict(load_ckpt['optimizer']) scheduler.load_state_dict(load_ckpt['scheduler'])
start_epoch = load_ckpt['epoch']+1 iter_epochs = range(start_epoch, args.epochs)
标签:load,ckpt,state,epoch,PyTorch,dict,断点,加载 From: https://www.cnblogs.com/booturbo/p/17358917.html