官方教程: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