체크포인트

체크포인트에 실제로 담아야 하는 것

체크포인트는 나중에 다시 불러올 수 있는 딕셔너리야. 모델 상태 사전, 옵티마이저 상태 사전, scheduler 상태 사전, GradScaler 상태 사전(혼합 정밀도를 쓴다면), 단계 번호, 에폭 번호, 설정을 담아야 해. 그래야 이 체크포인트로 무엇을 학습했는지 알 수 있어.

모델 상태 사전만 저장해도 추론에는 문제가 없지만, 학습을 재개하기에는 부족해. 옵티마이저와 scheduler에는 모델만으로 복원할 수 없는 이동 평균과 momentum이 들어 있거든.

팁: 옵티마이저도 항상 저장해. 누군가 '모델만' 불러와서 재개했다가 공들여 조정한 학습률이 0부터 다시 시작하는 모습을 정말 셀 수 없이 많이 봤어.

최적 체크포인트와 최신 체크포인트

최적 — 지금까지 검증 지표가 가장 높았던 체크포인트야. 배포할 때 사용해.

최신 — 지표와 관계없이 가장 최근에 저장한 체크포인트야. 중단 후 학습을 재개할 때 사용해.

둘 다 필요해. 최적은 운영을 위한 것이고, 최신은 엔지니어링을 위한 것이야. 서로 다른 파일명으로 저장해서 한쪽이 다른 쪽을 덮어쓰지하지 않게 해.

저장 위치

짧은 학습이라면 로컬 디스크에 저장해도 괜찮아. 여러 날에 걸친 실행이라면 노드가 죽어도 체크포인트가 남도록 S3 / GCS / 파일 서버에 저장해. 일부 학습 프레임워크(PyTorch Lightning, Composer, HF Trainer)는 이 작업을 대신 처리해 줘.

원칙: 체크포인트 저장은 보험이야. 학습이 80% 진행된 시점에 무언가 중단해서 GPU 시간 한 주를 날리지 않게 해 주는 가장 싼 일이야.

Code

전체 학습 상태 저장과 재개·python

import torch

def save_ckpt(path, model, opt, sch, scaler, step, epoch, cfg, metric):
    torch.save({
        "model": model.state_dict(),
        "optimizer": opt.state_dict(),
        "scheduler": sch.state_dict() if sch else None,
        "scaler": scaler.state_dict() if scaler else None,
        "step": step,
        "epoch": epoch,
        "config": cfg,
        "best_metric": metric,
    }, path)

def load_ckpt(path, model, opt=None, sch=None, scaler=None):
    ckpt = torch.load(path, map_location="cpu")
    model.load_state_dict(ckpt["model"])
    if opt and ckpt.get("optimizer"):
        opt.load_state_dict(ckpt["optimizer"])
    if sch and ckpt.get("scheduler"):
        sch.load_state_dict(ckpt["scheduler"])
    if scaler and ckpt.get("scaler"):
        scaler.load_state_dict(ckpt["scaler"])
    return ckpt["step"], ckpt["epoch"], ckpt.get("best_metric", 0.0)

Exercise

5에폭 동안 학습하고 체크포인트를 저장한 다음, 프로세스를 종료해. 체크포인트에서 재개해서 5에폭을 더 학습해. 검증 곡선은 재개 지점을 지나서도 연속으로 이어져야 해. 어느 방향으로도 jump가 생기면 안 돼. Jump가 있다면 저장하거나 불러오는 과정에서 무언가 빠진 거야.