모델 저장과 로드

학습 목표

이 레슨을 완료하면:

• •state_dict 방식과 전체 모델 저장 방식의 차이를 이해합니다
• •체크포인트를 사용해 학습을 중단하고 재개할 수 있습니다
• •GPU와 CPU 간 모델을 이동하는 방법을 익힙니다
• •실무에서의 모델 저장 전략을 설명할 수 있습니다

---

왜 모델 저장이 중요한가?

비유: 모델 저장은 게임의 "세이브" 기능과 같습니다. 몇 시간 동안 열심히 학습(플레이)한 결과를 저장하지 않으면, 프로그램을 종료하는 순간 모든 것이 사라집니다. 특히 GPU로 며칠간 학습하는 대규모 모델에서는 저장이 필수입니다.

모델을 저장해야 하는 상황:

---

3가지 저장 방식 비교

---

방식 1: state_dict 저장 (권장)

비유: state_dict 저장은 "레시피는 따로 보관하고, 양념 비율(가중치)만 저장"하는 것입니다. 레시피(모델 구조)가 있으면 양념 비율만으로 같은 요리를 재현할 수 있습니다.

state_dict란?

state_dict는 파이썬 딕셔너리로, 각 레이어의 이름을 키(key)로, 해당 레이어의 가중치 텐서를 값(value)으로 가집니다.

python
실행복사
# state_dict 내용 확인
model = MyModel()
print(model.state_dict().keys())
# 출력 예시:
# odict_keys(["fc1.weight", "fc1.bias", "fc2.weight", "fc2.bias"])

# 각 레이어의 가중치 크기 확인
for name, tensor in model.state_dict().items():
    print("{}: {}".format(name, tensor.shape))

저장하기

python
실행복사
# 모델의 가중치(state_dict)만 저장합니다
torch.save(model.state_dict(), "model_weights.pt")
# .pt 또는 .pth 확장자를 관례적으로 사용합니다

로드하기

python
실행복사
# 1단계: 모델 구조를 먼저 생성합니다 (코드가 필요!)
model = MyModel()

# 2단계: 저장된 가중치를 모델에 로드합니다
model.load_state_dict(torch.load("model_weights.pt"))

# 3단계: 추론(예측)에 사용할 경우 평가 모드로 전환합니다
model.eval()

왜 state_dict가 권장되는가?

1. •PyTorch 버전 호환성: 다른 버전에서도 로드 가능
2. •유연한 수정: 모델 구조를 변경한 후 일부 가중치만 로드 가능
3. •파일 크기: 가중치만 저장하므로 상대적으로 작음
4. •보안: 전체 모델 저장은 Pickle을 사용하므로 보안 위험 존재

---

방식 2: 전체 모델 저장

비유: 전체 모델 저장은 "완성된 요리를 통째로 냉동 보관"하는 것입니다. 꺼내면 바로 먹을 수 있지만, 냉동 방식(Pickle)에 의존하므로 해동이 안 될 수도 있습니다.

저장하기

python
실행복사
# 모델 전체를 저장합니다 (구조 + 가중치)
torch.save(model, "model_full.pt")

로드하기

python
실행복사
# 모델 클래스 코드 없이도 바로 사용 가능
model = torch.load("model_full.pt")
model.eval()

전체 모델 저장의 주의사항

이러한 이유로, PyTorch 공식 문서에서도 state_dict 방식을 권장합니다.

---

방식 3: 체크포인트 저장

비유: 체크포인트는 게임의 "풀 세이브"입니다. 캐릭터 상태(모델), 장비(옵티마이저), 진행도(에폭), 최고 점수(최적 성능) 등 학습을 이어가기 위한 모든 정보를 저장합니다.

왜 체크포인트가 필요한가?

모델 가중치만 저장하면 학습을 재개할 수 없습니다. 학습을 이어가려면:

• •모델 가중치 (어디까지 학습했는지)
• •옵티마이저 상태 (모멘텀 등 누적 정보)
• •현재 에폭 (어디서부터 이어갈지)
• •학습률 스케줄러 상태
• •최고 검증 성능 (비교 기준)

이 모든 것이 필요합니다.

체크포인트 저장하기

python
실행복사
def save_checkpoint(model, optimizer, epoch, val_loss, filepath):
    """학습 상태 전체를 체크포인트로 저장"""
    checkpoint = {
        "epoch": epoch,
        "model_state_dict": model.state_dict(),
        "optimizer_state_dict": optimizer.state_dict(),
        "val_loss": val_loss,
    }
    torch.save(checkpoint, filepath)
    print("Checkpoint saved: {}".format(filepath))

# 사용 예시: 매 에폭마다 또는 최고 성능일 때 저장
save_checkpoint(model, optimizer, epoch, val_loss, "checkpoint.pt")

체크포인트에서 학습 재개하기

python
실행복사
def load_checkpoint(filepath, model, optimizer):
    """체크포인트를 로드하고 학습 상태를 복원"""
    checkpoint = torch.load(filepath)

    # 모델 가중치 복원
    model.load_state_dict(checkpoint["model_state_dict"])

    # 옵티마이저 상태 복원 (모멘텀 등)
    optimizer.load_state_dict(checkpoint["optimizer_state_dict"])

    # 메타 정보 반환
    epoch = checkpoint["epoch"]
    val_loss = checkpoint["val_loss"]

    print("Resumed from epoch {}, val_loss: {:.4f}".format(epoch, val_loss))
    return epoch, val_loss

# 사용 예시
model = MyModel().to(device)
optimizer = optim.Adam(model.parameters(), lr=0.001)

start_epoch, best_val_loss = load_checkpoint(
    "checkpoint.pt", model, optimizer
)

# 이어서 학습
for epoch in range(start_epoch + 1, total_epochs):
    train_loss, train_acc = train_one_epoch(...)
    val_loss, val_acc = validate(...)
    # ... 나머지 학습 코드 ...

---

GPU와 CPU 간 모델 이동

실무에서는 GPU에서 학습하고 CPU에서 추론(서비스)하는 경우가 많습니다. 디바이스 간 모델 이동 방법을 알아야 합니다.

시나리오별 로드 방법

GPU에서 저장, CPU에서 로드

python
실행복사
# === GPU에서 학습 후 저장 ===
model = MyModel().cuda()
# ... 학습 진행 ...
torch.save(model.state_dict(), "model_gpu.pt")

# === CPU 환경에서 로드 ===
model = MyModel()  # CPU에 모델 생성
# map_location으로 텐서를 CPU로 매핑
model.load_state_dict(
    torch.load("model_gpu.pt", map_location=torch.device("cpu"))
)
model.eval()

CPU에서 저장, GPU에서 로드

python
실행복사
# === CPU에서 저장 ===
torch.save(model.state_dict(), "model_cpu.pt")

# === GPU 환경에서 로드 ===
model = MyModel()
model.load_state_dict(torch.load("model_cpu.pt"))
model = model.cuda()  # 모델을 GPU로 이동
model.eval()

왜 map_location이 필요한가?

GPU에서 저장한 텐서는 기본적으로 "cuda:0"이라는 위치 정보를 가지고 있습니다. CPU만 있는 환경에서 이 파일을 로드하면 "cuda:0 디바이스를 찾을 수 없다"는 오류가 발생합니다. map_location은 이 위치 정보를 재매핑합니다.

---

실무 저장 전략

전략 1: 최고 성능 모델만 저장

가장 기본적이고 많이 사용하는 전략입니다.

python
실행복사
best_val_loss = float("inf")

for epoch in range(epochs):
    train_loss, _ = train_one_epoch(model, train_loader, criterion, optimizer, device)
    val_loss, val_acc = validate(model, val_loader, criterion, device)

    if val_loss < best_val_loss:
        best_val_loss = val_loss
        torch.save(model.state_dict(), "best_model.pt")
        print("Epoch {}: New best model saved (val_loss={:.4f})".format(
            epoch + 1, val_loss))

전략 2: 주기적 체크포인트 + 최고 모델

긴 학습에서는 주기적 체크포인트와 최고 모델을 함께 저장합니다.

python
실행복사
for epoch in range(epochs):
    train_loss, _ = train_one_epoch(...)
    val_loss, val_acc = validate(...)

    # 최고 성능 모델 저장
    if val_loss < best_val_loss:
        best_val_loss = val_loss
        torch.save(model.state_dict(), "best_model.pt")

    # 10 에폭마다 체크포인트 저장 (학습 재개용)
    if (epoch + 1) % 10 == 0:
        save_checkpoint(model, optimizer, epoch, val_loss,
                       "checkpoint_epoch_{}.pt".format(epoch + 1))

전략 3: 마지막 N개 체크포인트만 유지

디스크 공간을 절약하면서도 안전하게 체크포인트를 관리합니다.

python
실행복사
import os

def save_with_rotation(model, optimizer, epoch, val_loss, max_keep=3):
    """최근 N개의 체크포인트만 유지"""
    filepath = "checkpoint_epoch_{}.pt".format(epoch + 1)
    save_checkpoint(model, optimizer, epoch, val_loss, filepath)

    # 오래된 체크포인트 삭제
    # 현재 에폭에서 max_keep 이전의 체크포인트를 찾아 삭제
    old_epoch = epoch + 1 - max_keep
    old_path = "checkpoint_epoch_{}.pt".format(old_epoch)
    if os.path.exists(old_path):
        os.remove(old_path)
        print("Removed old checkpoint: {}".format(old_path))

---

자주 발생하는 오류와 해결법

python
실행복사
# 부분 로드 (모델 구조가 약간 변경된 경우)
model.load_state_dict(
    torch.load("model.pt"),
    strict=False  # 일치하지 않는 키를 무시
)

---

핵심 요약

---

학습 체크리스트

• • state_dict와 전체 모델 저장의 차이를 설명할 수 있다
• • state_dict 방식이 권장되는 이유 4가지를 안다
• • 체크포인트에 포함해야 하는 정보를 나열할 수 있다
• • GPU에서 저장한 모델을 CPU에서 로드하는 방법을 안다
• • map_location이 필요한 이유를 설명할 수 있다
• • 실무에서의 저장 전략(최고 모델 + 주기적 체크포인트)을 이해한다