모델 저장과 불러오기

학습 결과는 모델 안의 숫자로 남는다

모델을 학습하면 Linear, Conv2d 같은 층 안의 가중치와 편향 값이 조금씩 바뀝니다. 이 값들이 바로 학습 결과입니다. 학습이 몇 분 만에 끝날 수도 있지만, 실제 프로젝트에서는 몇 시간 또는 며칠이 걸릴 수도 있습니다. 그래서 학습이 끝난 모델을 파일로 저장해 두어야 합니다.

PyTorch에서 가장 권장되는 방식은 모델 전체가 아니라 파라미터만 저장하는 것입니다. 이 파라미터 묶음을 state_dict라고 부릅니다.

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

state_dict에는 각 층의 가중치와 편향이 이름과 함께 저장됩니다. 예를 들어 fc1.weight, fc1.bias 같은 값들이 들어갑니다. 모델 구조 자체가 아니라, 학습된 숫자만 저장한다고 이해하면 됩니다.

불러올 때는 같은 모델 구조가 필요하다

state_dict는 파라미터만 담고 있으므로, 불러올 때는 먼저 같은 모델 클래스로 빈 모델을 다시 만들어야 합니다.

model = SimpleNet()
state_dict = torch.load("model_weights.pth", weights_only=True)
model.load_state_dict(state_dict)

여기서 중요한 점은 저장할 때 사용한 모델 구조와 불러올 때 만든 모델 구조가 같아야 한다는 것입니다. 층 이름이나 크기가 다르면 저장된 가중치를 어디에 넣어야 할지 알 수 없어서 오류가 납니다.

weights_only=True는 가중치를 불러오는 데 필요한 정보만 제한적으로 읽도록 하는 옵션입니다. 공식 튜토리얼에서도 모델 가중치를 불러올 때 이 방식을 좋은 습관으로 안내합니다.¹ 오래된 예제에서는 이 옵션이 빠져 있을 수 있지만, 새 코드를 작성할 때는 가능하면 함께 적어 두는 편이 안전합니다.

GPU에서 저장한 모델을 CPU 환경에서 불러올 수도 있습니다. 이때는 map_location을 지정하면 안전합니다.

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = SimpleNet()
state_dict = torch.load(
    "model_weights.pth",
    map_location=device,
    weights_only=True,
)
model.load_state_dict(state_dict)
model.to(device)

추론 전에는 eval 모드로 바꾼다

모델을 불러온 뒤 예측에 사용할 때는 model.eval()을 호출해야 합니다.

model.eval()

with torch.no_grad():
    outputs = model(inputs)
    predicted = outputs.argmax(dim=1)

eval()은 모델을 평가 모드로 바꿉니다. 특히 Dropout이나 BatchNorm 같은 층은 학습할 때와 추론할 때 동작이 다릅니다. 초보 단계에서는 이런 층을 직접 많이 쓰지 않더라도, 추론 전에는 항상 model.eval()을 호출하는 습관을 들이는 것이 좋습니다.

torch.no_grad()는 기울기 계산을 끄는 역할을 합니다. 예측만 할 때는 역전파가 필요 없으므로 메모리를 아끼고 실행도 더 가볍게 만들 수 있습니다.

모델 전체 저장보다 state_dict 저장을 권장한다

PyTorch에서는 다음처럼 모델 전체를 저장할 수도 있습니다.

torch.save(model, "full_model.pth")

이 방식은 간단해 보이지만, 초보자에게는 권장하지 않습니다. 모델 클래스가 정의된 위치나 코드 구조가 바뀌면 불러올 때 문제가 생길 수 있기 때문입니다. 모델 전체 저장은 내부적으로 Python의 직렬화 기능에 더 많이 의존하므로, 신뢰할 수 없는 파일을 불러오는 일도 피해야 합니다.

반면 state_dict 방식은 더 명확합니다.

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

모델 구조는 코드로 관리하고, 학습된 값은 파일로 관리합니다. 이 방식이 PyTorch에서 가장 흔히 쓰이는 저장 방식입니다.

가장 좋은 모델만 저장하기

학습을 여러 epoch 동안 진행하면 마지막 epoch의 모델이 항상 가장 좋은 것은 아닙니다. 검증 데이터 성능이 중간에 가장 좋았다가 이후에는 과적합으로 나빠질 수도 있습니다. 그래서 보통은 검증 성능이 좋아졌을 때만 저장합니다.

best_acc = 0.0

for epoch in range(num_epochs):
    train_loop(model, train_loader, loss_fn, optimizer)
    val_acc = evaluate(model, val_loader)

    if val_acc > best_acc:
        best_acc = val_acc
        torch.save(model.state_dict(), "best_model.pth")
        print(f"best model saved: {best_acc:.4f}")

이렇게 저장한 best_model.pth는 나중에 추론용 모델로 사용하기 좋습니다.

추가 학습까지 이어 가고 싶다면 모델 파라미터뿐 아니라 옵티마이저 상태와 epoch도 함께 저장할 수 있습니다.

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

다시 이어서 학습하려면 모델과 옵티마이저를 먼저 만든 뒤 체크포인트 딕셔너리를 불러와 각각 채웁니다.

checkpoint = torch.load(
    "checkpoint.pth",
    map_location=device,
    weights_only=True,
)

model.load_state_dict(checkpoint["model_state_dict"])
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
start_epoch = checkpoint["epoch"] + 1
best_acc = checkpoint["best_acc"]

정리하면, 예측용으로는 model.state_dict()를 저장하는 방식부터 익히면 충분합니다. 같은 모델 클래스를 다시 만들고, load_state_dict()로 값을 채운 뒤, eval() 모드에서 추론하면 됩니다.

참고

¹Save and Load the Model — PyTorch Tutorials