[pytorch] pytorch 모델 로드 중 Missing key(s) in state_dict 에러

pytorch로 딥러닝 모델을 저장하고 불러올 때 Missing key(s) in state_dict 라는 런타임 에러가 발생하는 경우가 종종 있다.

 

대부분 딥러닝 모델과 불러온 모델 웨이트의 키값이 맞지 않아서 발생하는 오류인데, 모델과 모델 웨이트가 서로 짝이 아닌 경우에 발생하고 간혹 그렇지 않은 경우에도 해당 에러가 발생해서 우리를 괴롭힌다... ㅎ

 

그 때 아래와 같이 torch.load로 모델의 state dict를 불러오고 디버깅을 해서 state dict의 키값을 확인해보면 정확하게 문제를 파악할 수 있다.

model_path = './model.pth'
model_state_dict = torch.load(model_path)

 

 

위 첫 번째 캡처를 보면 모델 state dict에 모델 웨이트와 학습 정보인 iter, optimizer, scaler 등의 키값이 함께 저장되어 있는 것을 볼 수 있다. 그래서 모델에 바로 불러오면 Missing key(s) 에러가 발생한다. 이 때 모델의 웨이트에 해당하는 키값으로 인덱싱해주고 모델 state dict를 불러오면 정상적으로 모든 키가 매치된다.

 

꼭 이런 경우가 아니더라도 불러온 모델 state dict의 키값들을 확인하면 완전히 다른 모델의 웨이트를 가져왔는지, 키값만 조금 다른 부분이 있는지 등의 문제를 파악할 수 있다.