KeyError: 'unexpected key "module.**" in state_dict'

발생한 에러 메세지:

KeyError: 'unexpected key "module.**" in state_dict'

 

모델 state_dict의 key 값들을 출력해보면, "module."으로 시작하는 key는 존재하지 않는 것을 확인할 수 있다.

그런데 왜 이런 key값으로 저장되었을까?

 

이유는 간단하다. 처음 모델의 state_dict를 저장할 때, 모델이 nn.DataParallel로 선언이 되어있을 경우 이런 문제가 발생한다. 따라서 문제를 해결하는 방법은 두 가지가 있다.

 

1. 로딩하는 모델을 nn.DataParallel 로 변환하기

state_dict를 저장할 때와 똑같이 모델을 nn.DataParallel로 변환한 뒤, state_dict 를 로딩한다.

GPU_IDs = [0,1,2,3] # DataParallel에 사용할 GPU ID들을 입력
model = torch.nn.DataParallel(model, device_ids = GPU_IDs) # DataParallel 선언
model.load_state_dict(state_dict) # state_dict 로딩

 

2. state_dict의 key값을 "module."을 없앤 값으로 재설정하기

Key가 일치하지 않기 때문에 발생하는 문제이므로 state_dict의 key를 직접 해당 모델의 key와 일치하도록 바꾸는 방법도 있다. 예를 들면, 기존의 key값이 "module.encoder.fc.weight"였다면, 이를 "encoder.fc.weight"로 바꿔주면 문제가 해결된다.

from collections import OrderedDict
new_state_dict = OrderedDict()
for k, v in state_dict.items():
name = k[7:] # "module."을 지우기
new_state_dict[name] = v
model.load_state_dict(new_state_dict) # "module."이 지워진 새로운 state_dict 로딩