파이토치 모델 저장

파이토치 모델 저장은 주로 torch.save(model.state_dict(), PATH)를 사용하여
가중치(state_dict)만 저장하는 방식이 권장된다.
이 방식은 용량이 작고 유연하며, 추론 시 model.load_state_dict()로 불러온다.

 

모델 전체를 저장(torch.save(model, PATH))할 수도 있으나,
클래스 구조 의존성 때문에 권장되지 않는다. 

 

1. State Dict 저장 (가중치만)
가장 권장되는 방식으로,
파라미터만 저장하므로 모델 클래스 구조가 정의되어 있어야 한다. 

• 저장: torch.save(model.state_dict(), 'model.pth')
• 불러오기:

  model = TheModelClass() # 모델 클래스 정의 필요
  model.load_state_dict(torch.load('model.pth'))
  model.eval() # 추론 전 필수

   

2. 전체 모델 저장 (구조 + 가중치)
모델 전체를 저장하여, 불러올 때 모델 클래스 정의가 필요 없다.

단, 프로젝트 파일 구조가 바뀌면 로드되지 않을 수 있다. 

• 저장: torch.save(model, 'model.pt')
• 불러오기: model = torch.load('model.pt')
   

3. 체크포인트 저장 (학습 재개용)
학습 중간에 가중치, 옵티마이저 상태, 에포크 등을 함께 저장한다. 

torch.save({
    'epoch': epoch,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'loss': loss,
}, 'checkpoint.pth')

 

4. 핵심 팁

• 확장자: 보통 .pt 또는 .pth를 사용한다.
• GPU/CPU: GPU에서 학습한 모델을 CPU에서 불러올 경우 map_location 설정을 사용한다.

  device = torch.device('cpu')
  model.load_state_dict(torch.load('model.pth', map_location=device))