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

Posted 2023-04-27 BooTurbo

 

 

 

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

一、神经网络模型的保存，基本两种方式：
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)

 

YOLOV5-断点训练/继续训练

YOLOV5-断点训练/继续训练

情况：
1.训练过程中中断了，继续训练
2.训练完了，但是未收敛，在这个基础上，还想用这个权重、学习率等参数继续训练多一些批次

1. 训练过程中中断了，继续训练

断电、或者什么原因中断了，比如设定epoch为200，但是在90这里中断了，想从断点这里继续 训练到 200个epoch
直接在trainds.py设置参数–resume的defualt=True即可（在代码 517行附近）

1.parser.add_argument('--resume', nargs='?', const=True, default=True, help='resume most recent training')

2.训练完原有epoch，发现没有收敛，想继续训练

训练完原有epoch后，但还继续训练，比如设置epoch为30，已经训练完了，但是没有收敛，想使用训练了30 epoch的权重继续训练40个epoch， 总共就是70个epoch

**注意：**不是直接在train.py设置weight参数为训练好的权重（last.pt），然后设置epoch为70，这样不行，因为学习率等超参数启动时还是按照默认的。（自己把一些超参数设置好也可以，比较麻烦）
在train.py代码里面设置就可以使用原有30 epoch的权重后继续训练到40epoch

step1：

先将–resume的defualt改为True，再将–epochs设为想要继续训练到的epochs，即defualt=70（在trainds.py代码的513行附近）

parser.add_argument('--epochs', type=int, default=70)

step2：

在trainds.py文件中，按下Ctrl+F搜索Epochs定位到大概197行的位置，做以下两处修改：

添加一行代码：ckpt[‘epoch’] = 30
修改代码：start_epoch = ckpt[‘epoch’]（后面不加1）
注意：这两处修改是为了断点训练，如果不再需要断点训练，请将这两处修改都还原。

step3：

在trainds.py文件中，按下Ctrl+F搜索replace定位到大概565行的位置，做以下两处修改：

如图添加一行代码：init_epochs = opt.epochs（记录最开始由epochs参数设定的数值70，因为replace这行代码，将epochs替换成了原始的epochs数值30）
修改代码：如图进行添加，reinstate恢复成设定的epochs

做完以上3个修改，训练就会从第30个epoch的基础上进行继续训练到70个epoch，和直接设置epoch为70一样

参考文献：https://blog.csdn.net/LWD19981223