Pytorch模型训练恢复之load_state_dict

官方教程：https://pytorch.org/tutorials/beginner/saving_loading_models.html#
新版pytorch中load_state_dict()不支持map_location参数，但torch.load()支持

背景

假设pytorch model执行save操作时，模型参数是在gpu上，那么运行checkpoint = torch.load(checkpoint_path)之后，checkpoint['backbone'], checkpoint['optimizer']中的参数，看其device属性，仍在gpu上，但是经过load_state_dict，比如backbone.load_state_dict(checkpoint['backbone'])，此时参数均在cpu上，接着backbone.cuda()可以将参数转换到gpu上。

问题

此时坑出现了，optimizer['param_groups']会随着backbone, head的转换而自动转换到gpu上，但是optimizer['state']仍然留在cpu上，所以运行optimizer.step()时，在这一步会报错：buf.mul_(momentum).add_(1 - dampening, d_p)，因为buf在cpu上，而梯度d_p是在gpu上。

解决

load_state_dict后手动将optimizer.state中的参数转移至gpu上，代码如下：

for k, v in optimizer.state.items():    # key is Parameter, val is a dict {key='momentum_buffer':tensor(...)}
    if 'momentum_buffer' not in v：
        continue
    optimizer.state[k]['momentum_buffer'] = optimizer.state[k]['momentum_buffer'].cuda()

其他

对于lr_scheduler，可以保存其_step_count, last_epoch这些参数，恢复训练时，可以基于恢复的optimizer重定义lr_scheduler（若训练中断前保存了整个lr_scheduler，恢复时这个scheduler是用不了的，因为它指向原来的optimizer），然后用上述保存的参数更新lr_scheduler中的对应参数。示例代码如下：

	checkpoint = torch.load(checkpoint_path)
	...
	optimizer.load_state_dict(checkpoint['optimizer'])
	lr_scheduler_saved = checkpoint['lr_scheduler']
	lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer)
	lr_scheduler._step_count = lr_scheduler_saved._step_count
	lr_scheduler.last_epoch = lr_scheduler_saved.last_epoch