[PYTORCH] 모델 전체 저장 vs 가중치만 저장의 3가지 결정적 차이와 권장 방법 및 해결 전략

모델 전체 저장 vs 가중치만 저장

 

PyTorch 직렬화의 심층 분석: 왜 전문가들은 state_dict를 고집하는가?

---

1. 서론: 모델 저장 방식의 선택이 프로젝트의 성패를 가른다

딥러닝 모델 학습은 수 시간에서 수개월이 걸리는 고된 작업입니다. 공들여 학습시킨 모델을 파일로 저장하는 방식에는 크게 두 가지가 있습니다. 바로 '모델 객체 전체(Entire Model)'를 저장하는 방식과 '가중치(state_dict)'만을 저장하는 방식입니다. 초보 개발자들은 편리함 때문에 전자를 선택하곤 하지만, 실무 환경이나 모델 배포 단계에서는 후자가 압도적으로 권장됩니다. 본 포스팅에서는 이 두 방식의 기술적 차이점과 발생 가능한 문제점, 그리고 실무에서 즉시 활용 가능한 7가지 솔루션을 상세히 다룹니다.

2. 모델 전체 저장 vs 가중치 저장 상세 비교

두 방식의 본질적인 차이는 'Python의 Pickle'이 데이터를 처리하는 범위에 있습니다. 전체 저장은 클래스 구조까지 직렬화에 포함시키며, 가중치 저장은 오직 숫자로 구성된 파라미터 텐서만을 저장합니다.

항목	모델 전체 저장 (Entire Model)	가중치만 저장 (state_dict)
저장 대상	모델 클래스 구조 + 가중치 파라미터	OrderedDict 형태의 가중치 텐서만
코드 의존성	매우 높음 (동일한 디렉토리 구조 필요)	낮음 (모델 정의 코드만 있으면 로드 가능)
유연성	낮음 (다른 프로젝트에서 활용 시 오류 잦음)	매우 높음 (부분 로딩, 전이 학습 용이)
안정성	버전 업데이트 시 깨질 확률 높음	매우 안정적이며 공식 권장 방식
불러오기 방식	torch.load(PATH)	model.load_state_dict(torch.load(PATH))

3. 왜 state_dict 방식이 '해결'의 열쇠인가?

전체 모델 저장은 저장 당시의 디렉토리 구조와 클래스 정의를 피클(Pickle) 파일 내부에 묶어버립니다. 따라서 파일을 다른 폴더로 옮기거나, 클래스 이름을 변경하기만 해도 AttributeError 또는 ModuleNotFoundError가 발생합니다. 반면 state_dict는 데이터(텐서)만 담고 있기 때문에, 모델의 아키텍처 코드가 존재하기만 한다면 언제 어디서든 안전하게 가중치를 주입할 수 있습니다.

---

4. 실무자를 위한 7가지 구체적 Example 및 해결 전략

Example 1: 공식 권장 방식(state_dict)의 표준 구현

가장 기본이 되면서도 가장 강력한 저장 및 로드 방법입니다.

# 1. 저장하기 (Best Practice)
torch.save(model.state_dict(), 'best_model.pt')

# 2. 로드하기 (아키텍처 선언 후 가중치 주입)
new_model = MyNet()
new_model.load_state_dict(torch.load('best_model.pt'))
new_model.eval() # 추론 모드 전환 필수!
        

Example 2: 클래스 구조 변경 시 발생하는 로드 오류 해결

모델의 일부 레이어를 수정했을 때, 일치하는 가중치만 불러오는 strict=False 기법입니다.

# 레이어 이름이 일부 바뀌었거나 구조가 수정되었을 때
# 일치하지 않는 키는 무시하고 로드
model.load_state_dict(torch.load('old_model.pt'), strict=False)
        

Example 3: 대규모 협업을 위한 체크포인트(Checkpoint) 구성

단순 가중치 외에 학습 상태를 모두 저장하여 '재현성'을 확보하는 방법입니다.

torch.save({
    'epoch': 150,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'scheduler_state_dict': scheduler.state_dict(),
    'best_loss': 0.001
}, 'checkpoint_v1.tar')
        

Example 4: 특정 GPU 번호에서 학습된 모델을 다른 장치로 로드

CUDA 장치 불일치 문제를 map_location으로 해결하는 방법입니다.

# GPU 0번에서 저장된 파일을 CPU로 불러올 때
state_dict = torch.load('model.pt', map_location=torch.device('cpu'))

# 특정 GPU 번호(예: 1번)로 직접 매핑할 때
state_dict = torch.load('model.pt', map_location={'cuda:0': 'cuda:1'})
        

Example 5: 데이터 병렬화(DataParallel) 모델 로드 해결 방법

nn.DataParallel로 저장하면 키 값에 module.이 붙는 문제를 해결하는 코드입니다.

state_dict = torch.load('dp_model.pt')
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)
        

Example 6: 전이 학습(Transfer Learning) 가중치 부분 이식

기존 모델의 가중치 중 필요한 부분만 추출하여 새 모델에 적용하는 실무 테크닉입니다.

pretrained_dict = torch.load('backbone.pt')
model_dict = model.state_dict()

# 현재 모델 아키텍처에 존재하는 키만 필터링
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}

# 가중치 업데이트 후 로드
model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)
        

Example 7: 가중치 파일 용량 최적화 (FP16/BF16 저장)

저장 용량을 50% 절감하면서도 배포 효율을 높이는 하프 프리시전(Half-precision) 저장법입니다.

# 모델을 반정밀도로 변환 후 가중치 추출
model.half() 
torch.save(model.state_dict(), 'model_fp16.pt')
        

5. 결론: 실무적인 선택 가이드

실험 단계의 빠른 백업이 목적이라면 전체 저장이 달콤하게 느껴질 수 있습니다. 하지만 배포, 공유, 장기적인 유지보수를 생각한다면 반드시 state_dict 방식을 사용하십시오. 특히 torch.save 시에는 파일 확장자로 .pt를 사용하여 PyTorch 관련 객체임을 명시하는 것이 관례입니다. 또한 로드 후에는 반드시 model.eval()을 호출하여 Dropout이나 BatchNorm 레이어가 추론 모드로 작동하도록 설정하는 것을 잊지 마십시오.

6. 내용 출처

• PyTorch Tutorials: "Saving and Loading Models" (https://pytorch.org/tutorials)
• PyTorch Discussion Forum: "Serialization of Models and Tensors"
• Hands-On Machine Learning with Scikit-Learn, Keras, and TensorFlow (2nd Edition)