PyTorch 모델 저장 시 ModuleNotFoundError 오류 해결기: pickle과 버전 호환의 함정

모델을 저장하고 다시 불러오기만 했을 뿐인데 에러가 났다. 그 원인은 다름 아닌 pickle이었다.

---

🧩 문제는 환경의 “미묘한 차이”였다

Stable Diffusion 기반 API 서버를 개발하는 과정에서 .pth 파일로 저장한 모델을 불러오는 도중 다음과 같은 오류를 만났다.

ModuleNotFoundError: No module named 'diffusers.models.unets'

 

흥미로운 점은 같은 서버에서 저장한 모델임에도, Docker 컨테이너에서 실행했을 때 이 오류가 발생했다는 것이다. 알고 보니 로컬 환경과 컨테이너 내부의 라이브러리 버전이 미묘하게 달랐고, 바로 그 차이로 인해 모델을 제대로 불러올 수 없었던 것이다.

---

왜 이런 일이 발생했을까?

# 저장 시
[UNet2DConditionModel 객체] 
     └── torch.save() → .pth 파일 (pickle 직렬화)

# 로드 시
torch.load() 
     └── diffusers.models.unets.UNet2DConditionModel 
             ✖️ 모듈 없으면 → ModuleNotFoundError

PyTorch에서 사용하는 torch.save()는 내부적으로 Python의 pickle을 사용해 객체를 직렬화한다. 이때 객체의 클래스 경로 정보(예: diffusers.models.unets.UNet2DConditionModel)가 문자열로 저장되기 때문에, 모델을 로드할 시점에 해당 경로가 정확히 존재해야만 오류 없이 불러올 수 있다.

즉, Docker 이미지가 생성되는 시점에 설치된 diffusers 패키지 구조가 로컬과 달라졌다면, torch.load()는 해당 클래스를 찾지 못하고 ModuleNotFoundError를 발생시킨다.

---

해결 방법

방법 1: 저장 당시와 동일한 라이브러리 버전 맞추기

[pip freeze > requirements.txt] → Docker 이미지 생성
       │
       ▼
[torch.load()] → 동일 구조 & 버전이면 정상 작동 ✅

가장 안정적인 방법은 모델을 저장한 환경과 완전히 동일한 버전의 라이브러리를 사용하는 것이다. 특히 다음 3가지 패키지는 반드시 버전을 고정하는 걸 추천한다:

diffusers==0.34.0 transformers==4.54.1 torch==2.2.2

이 버전들을 requirements.txt에 명시해 Docker 이미지를 구성하면, 모델 로딩 시 문제가 생기지 않는다.

💡 팁: 모델 저장 당시의 pip freeze를 기록해두면 나중에 환경 재현이 훨씬 쉬워진다.

---

방법 2 (참고용): state_dict만 저장

pickle 직렬화 자체를 피하고 싶다면, state_dict() 방식으로 모델의 파라미터만 저장할 수 있다:

# 저장 
torch.save(model.state_dict(), "unet_state_dict.pth") 

# 로드 
model = UNet2DConditionModel(...) # 구조 명시 필요 
model.load_state_dict(torch.load("unet_state_dict.pth"))

 

이 방식은 라이브러리 버전이 달라도 잘 작동하지만, 모델 구조를 직접 다시 정의해야 해서 번거롭다. 라이브러리 구조가 자주 바뀌지 않는 프로젝트라면 고려해볼 만한 대안이다.

---

정리: Docker 환경에서는 버전 관리가 생명

• torch.save()와 torch.load()는 pickle 기반이라 구조 의존성이 매우 크다.
• Docker 환경을 새로 빌드하거나 서버 환경이 바뀔 때마다 반드시 다음 세 가지 버전을 체크하자:
  • diffusers
  • transformers
  • torch
• ModuleNotFoundError가 떴다면, 가장 먼저 라이브러리 버전 불일치를 의심해야 한다.

---

후속 작업으로 남긴 것들

이번 경험을 바탕으로 앞으로는 다음을 실천하려 한다:

• 모델 저장 시점의 requirements.txt를 반드시 백업해두기
• 모델을 배포할 컨테이너에는 정확한 라이브러리 버전 명시
• 장기적으로는 state_dict 방식도 병행 검토

---

이 문제는 작지만 은근히 자주 발생하고, 특히 협업 중에 환경이 미묘하게 달라졌을 때 디버깅 시간을 잡아먹는다. 나처럼 이 문제에 막히는 개발자 분들께 이 글이 작은 힌트가 되길 바란다!