Pytorch：训练中断再恢复时的注意事项

参考链接：Pytorch：lr_schedule恢复训练的注意事项

import torch
import torch.optim as optim
from torchvision.models.resnet import resnet50
net = resnet50(num_classes=1000)
optimizer = optim.Adam(net.parameters(), lr=1e-3)
scheduler = optim.lr_scheduler.MultiStepLR(optimizer, [20, 30, 40, 50], 0.1)
for epoch in range(num_epoch):
    scheduler.step()
    train()
    valid()
    ...

 

      有时候会因为不可抗拒的外界因素导致训练被中断，在pytorch中恢复训练的方法就是把最近保存的模型重新加载，然后重新训练即可。假设我们从epoch10开始恢复训练，可以利用lr_scheduler的last_epoch参数,last_epoch指定为10，则网络会从第11个epoch开始接着训练。

scheduler = optim.lr_scheduler.MultiStepLR(optimizer, [20, 30, 40, 50], 0.1, last_epoch=10)

需要注意的是在optimizer定义参数组的时候需要加入’initial_lr’，不然会报错：

KeyError: "param 'initial_lr' is not specified in param_groups[0] when resuming an optimizer" 

这个时候需要将代码修改为如下形式，其中的‘initial_lr'就等于我们设置的初始化学习率，下面的循环也要修改为开始训练的epoch。

import torch
import torch.optim as optim
from torchvision.models.resnet import resnet50
net = torch.load('resnet50_epoch10.pth')
optimizer = optim.Adam([{'params': net.parameters(), 'initial_lr': 1e-3}], lr=1e-3)
scheduler = optim.lr_scheduler.MultiStepLR(optimizer, [20, 30, 40, 50], 0.1, last_epoch=10)
for epoch in range(11, num_epoch):
    scheduler.step()
    train()
    valid()
    ...