[PyTorch] torch.load error

(feat. ModuleNotFoundError: No module named 'models')

[error] torch.load error (feat. ModuleNotFoundError: No module named 'models')

상황

 

직접 만든 모델을 학습 후

torch.save(model, path)

로 저장한 후

torch.load(weight)

을 하는 도중 다음과 같은 에러가 발생하였다.

 

ModuleNotFoundError: No module named 'models'

 

원인

 

라이브러리를 통해 model을 학습 후 저장한 것이 아닌 직접 만든 모델이나 git clone을 통해 받은 모델을 학습시킨 후 torch.save()를 통해 weight을 저장한 후 model이 담겨있는 py파일 경로가 아닌 다른 경로에서 load하는 경우 발생하는 에러이다.

 

 

해결

 

import sys

sys.path.insert(0, yourmodelpath)

 

위의 코드와 같이 sys.path에 모델의 파일이 담긴 경로를 추가해주어 해결할 수 있다.