在对模型进行微调训练的时候,代码报错如下:
File "/home/xxxx/EdgeNets/utilities/train_eval_seg.py", line 48, in train_seg
optimizer.step()
File "/home/xxxx/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/optim/sgd.py", line 100, in step
buf.mul_(momentum).add_(1 - dampening, d_p)
RuntimeError: expected device cpu but got device cuda:0
解决的方法是,重载optimizer的参数时将所有的tensor都放到cuda上(加载时默认放在cpu上了),代码片段如下:
checkpoint = torch.load(filename)
net.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
for state in optimizer.state.values():
for k, v in state.items():
if torch.is_tensor(v):
state[k] = v.cuda()