[이거 어떡하지] optim in step: expected device cpu but got device cuda:0

증상:

optim.step() - RuntimeError: expected device cpu but got device cuda:0

1. pytorch에서 코딩하는데 난 model도 cuda에 올렸고 (Model().cuda() 혹은 Model().to(device)를 사용)

2. torch.load({CHECKPOINT_PATH})로 optimzer의 state를 불러왔는데

 

아래와 같이 optimizer의 step()에서 이러한 증상이 나타납니다. 왜 그럴까요?

optimizer에서 device가 다르다고 에러가 뜨기도 하나?

 

해결 방법:

1] optim_state = torch.load({CHECKPOINT_PATH})["optimizer"]
2] optim.load_state_dict(optim_state)
3] for state in optim.state.values():
4]     for k, v in state.items():
5]         if torch.is_tensor(v):
6]             state[k] = v.to(device)

1]
torch.load함수로 checkpoint를 불러와서 optimzer key에 있는 optimizer의 state를 불러옵니다. (torch.save에서 optimizer를 저장한 key 이름대로 불러오시면 됩니다.)
2]
optim이라는 optimizer의 state를 update합니다.(optimizer중 Adam의 경우 model parameter마다 할당된 momentum과 adaptive variable을 update하게 됩니다.)
3]~6] 
device가 "cuda" device로 설정되어있을시, optim의 state마다 cuda에 올려줍니다. model은 GPU에 올라가있지만 optimizer의 state가 올라가있지 않기 때문에 optimizer는 cpu tensor가 오길 바라지만 cuda tensor가 오게 되어 위와같은 에러를 발생시킵니다.

 

 

위와 같이 optimizer state를 gpu로 올려주면 더이상 optimizer.step()에서 에러가 발생하지 않습니다.

 

해결될 줄 알았는데 안되는 방법:

1) optim.to(device)는 먹히지 않습니다. optimizer class는 to라는 메소드를 갖고있지 않기 때문입니다.

2) torch.load() 함수의 map_location=device argument도 먹히지 않는 것을 확인하였습니다. 혹시 이것으로 해결하신 분은 알려주시면 감사하겠습니다.