[AI] PyTorch 파일 확장자 (.pt, .pth, .ckpt)

모델 파일이란?

 PyTorch에서는 학습을 통해서 생성된 모델이나 학습 진행 상황을 저장하기 위해서 주로 `.pt` , `.pth` 두 가지 파일 확장자를 사용한다.

저장된 모델 파일을 활용해서 재학습을 생략하고 재사용성을 높일 수 있다.

이렇게 향상된 재사용성을 통해서 추론이나 추가적인 훈련 같은 후속 작업을 할 때 시간적, 자원적으로 절약할 수 있다.

💡모델 파일
학습을 통해서 생성된 모델의 구조와 훈련된 매개변수, 학습된 정보 등을 담고 있는 파일

 

‘.pt’, ‘.pth'그리고 추가로 ‘.ckpt’는 일반적으로 PyTorch 같은 기계학습 프레임워크에서 모델, 가중치, 체크포인트 등을 저장하는데 사용된다.

기능적으로 유사하지만, 사용법과 규칙, 관행에 따라서 다르게 쓰이는 경우가 많다.

따라서, 사용할 때 정확한 규칙과 이유를 알고 사용해야 한다.

 

재교육 대신 모델을 저장하는 이유

• 비용 효율적
  • 재교육에는 시간, GPU 비용 등 많은 컴퓨팅 리소스가 소비된다.
  • 규모가 큰 모델일수록 재교육에 소비되는 자원이 매우 많아진다.
• 재현성
  • 저장한 모델을 재사용할 때 일관된 결과를 보장한다.
  • 일정한 결과를 요구하는 연구나 생산과정에서 매우 중요하다.
• 강건성
  • 예상치 못한 실패로 인해서 작업이 손실되는 것을 방지한다.
  • 오류로 인해서 멈춘 교육 과정을 재교육과 손실 없이 다시 재개할 수 있다.

=> 본질적으로 기계 학습 과정에서의 효율성, 안정성 및 원할한 공동작업을 위해서 매우 중요하다.

.pt

모델의 전체를 저장한다.

📌모델의 전체
- 구조(Structure)
- 가중치(Weight)
- 파라미터(Parameters)
      - optimizer
      - loss
      - epoch 등등

 

목적

• 배포를 목적으로 모델의 전체를 저장한다.
• 추론을 위해서 모델의 전체를 저장한다.
• TorchScript 기반의 모델을 저장한다.

장점

• 유연하게 독립적인 모델의 전체를 저장할 수 있다.
• 간단하게 배포할 수 있다.
• 추론 준비 모델에 적합하다.

단점

• 모델의 전체를 저장하기 때문에 파일의 용량이 커진다.
• 훈련을 재개할 때 추가적인 작업이 필요하다.

규칙

• 생산 및 배포를 위해서 준비된 독립적인 모델 전체를 저장할 때 사용한다.

예시 코드

📥  모델의 전체 저장

torch.save(model, "model.pt")

 

📤 모델의 전체 로드

model = torch.load("model.pt")
model.eval()    #모델을 평가모드(evaluation mode)로 설정한다.

 

.pth

모델의 가중치만 저장한다.

 

목적

• 모델의 상태 사전(state_dict) 저장에 사용된다.
• 학습 중이거나, 아키텍처 없이 모델 가중치를 저장할 때 사용된다.

💡상태 사전(state_dict)
모델의 각 계층을 매개변수 텐서로 매핑되는 Python 사전(dict) 객체

 

장점

• 훈련을 재개하거나, 가중치를 다른 아키텍처로 이전하는 데 유연하다.
• 파일 용량이 작다.

단점

• 가중치만 저장되기 때문에 수동적으로 메타데이터를 관리해줘야 한다.
• python path와 충돌 위험이 있다.
• TorchScript 에서 사용 불가능하다.

규칙

• 훈련중인 모델의 가중치를 저장할 때 사용된다.

예시 코드

📥  모델의 가중치(state_dict) 저장

torch.save(model.state_dict(), "model_weights.pth")

 

📤 모델에 가중치(state_dict) 로드

model.load_state_dict(torch.load("model_weights.pt"))
model.eval()

 

.pt VS .pth

PyTorch 에서 사용되는 ‘.pt’ 와 ‘.pth’ 파일 확장자 사이에는 별다른 차이점이 없다.

파일 형식이나 내용 면에서 일치하고, 상호교환 가능하다.

 

즉, 본질적으로 두 파일 확장자는 차이가 크게 없으며, 어느 확장자를 사용하더라도 PyTorch 모델을 저장할 수 있다.

코드상으로도 두 파일 확장자 상관없이 동일하게 사용할 수 있다.

 

하지만, 규칙과 관행에 따라서 다르게 저장하고, 다르게 사용된다.

.pt : 모델의 전체 저장
.pth : 모델의 가중치만 저장

 

.ckpt

모델의 체크포인트를 저장한다.

📌체크포인트
- 모델 가중치
- 최적화 상태
- 스케줄러 상태
- 교육 메타데이터 등

 

목적

• 교육, 특히 다중 프레임워크 시나리오 또는 PyTroch Lightning과 같은 도구에서 체크포인트를 저장하는데 사용된다.

장점

• 모델의 가중치(state_dict) 뿐만 아니라 다양한 메타데이터를 저장할 수 있다.
• 중단된 훈련을 재개할 때, 정확한 지점에서 원활하게 재개할 수 있다.
• PyTroch Lightning에서 지원이 내장되어 있다.

단점

• 가중치 외에 추가적으로 데이터가 저장되기 때문에, 파일 용량이 크다.
• 배포용으로는 부적합하다.
• 호환성이 좋지 않다.

규칙

• 진행 상황과 메타데이터를 보존해야 하는 체계적인 학습 파이프라인에서 더 일반적으로 사용된다.

예시 코드(Torch)

📥  모델의 체크포인트 저장(수동)

checkpoint = {
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'epoch': epoch,
}
torch.save(checkpoint, "model.ckpt")

 

📤 모델에 체크포인트 로드(수동)

checkpoint = torch.load("model.ckpt")
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']

 

예시 코드(PyTorch Lightining)

📥  체크포인트 저장

from pytorch_lightning import Trainer
from my_model import MyModel # 사용하고자 하는 모델

trainer.save_checkpoint("model.ckpt")

 

📤 체크포인트 로드

model = MyModel.load_from_checkpoint("model.ckpt")

 

정리

• .pt
  • 모델의 전체를 저장할 때 사용한다.
  • 모델을 배포할 때 사용한다.
  • 추론 준비 모델에 사용한다.
• .pth
  • 모델의 가중치를 저장할 때 사용한다.
  • 학습, 실험에서 유연하게 여러 번 재사용할 때 사용한다.
• .ckpt
  • 모델의 체크포인트를 저장할 때 사용한다.
  • 전체적인 훈련을 멈췄다가 재개하는 상황에 사용한다.