最近在学习使用pytorch,遇到一个坑。
从中断的训练中继续训练需要重新加载检查点,一些优化器(比如adam)的一些变量也需要被保存到检查点中。
而在使用load_state_dict()还原时,有可能会将这些变量还原到CPU上,如下面的代码这样:
model = Model()
optimizer = Optim()
checkpoint = torch.load(ckp_path)
model.load_state_dict(checkpoint["model"])
optimizer.load_state_dict(checkpoint["optimizer"])
model.cuda().train()
......
像这样的话,优化器的变量就会被还原到CPU上,这在训练更新权重时会导致错误,因为模型权重在GPU上,而优化器的变量在CPU上。
造成这个问题的原因是在还原检查点之前,模型没有被转到GPU上,这样还原时,所有模型相关的变量都会被还原到CPU上。
解决办法有三。
方法一:手动转换
model = Model()
optimizer = Optim()
checkpoint = torch.load(ckp_path)
model.load_state_dict(checkpoint["model"])
optimizer.load_state_dict(checkpoint["optimizer"])
for state in optimizer.state.values():
for k, v in state.items():
if isinstance(v, torch.Tensor):
state[k] = v.cuda()
model.cuda().train()
......
方法二:还原检查点之前,将模型转到GPU
model = Model().cuda()
optimizer = Optim()
checkpoint = torch.load(ckp_path)
model.load_state_dict(checkpoint["model"])
optimizer.load_state_dict(checkpoint["optimizer"])
model.cuda().train()
......
方法三:在load中设置map_location
device = torch.device("cuda")
model = Model()
optimizer = Optim()
checkpoint = torch.load(ckp_path, map_location=device)
model.load_state_dict(checkpoint["model"])
optimizer.load_state_dict(checkpoint["optimizer"])
model.cuda().train()
......