학습이 끝났다면, 이제 결과를 붙잡아야 한다
지난 장에서 우리는 학습 루프를 만들었다. 배치를 꺼내고, 모델이 예측하고, 손실을 계산하고, loss.backward()로 기울기를 구하고, optimizer.step()으로 가중치를 조금씩 고쳤다. 이 반복이 충분히 진행되면 모델은 처음보다 훨씬 나은 예측을 하게 된다.
그런데 여기서 중요한 질문이 생긴다.
학습이 끝난 모델은 어디에 남아 있을까?
코드가 실행되는 동안에는 모델이 메모리 안에 있다. 하지만 노트북을 닫거나 프로그램이 종료되면, 힘들게 학습한 가중치도 함께 사라진다. 30분 동안 학습한 모델이라면 다시 학습하면 된다. 하지만 10시간, 3일, 몇 주 동안 학습한 모델이라면 이야기가 달라진다. 저장하지 않은 모델은 완성된 그림을 저장하지 않고 그림판을 닫아 버린 것과 같다.
딥러닝에서 모델 저장은 단순한 부가 기능이 아니다. 실습 결과를 재사용 가능한 결과물로 바꾸는 마지막 단계다. 저장할 수 있어야 나중에 추론에 사용할 수 있고, 다른 사람에게 전달할 수 있으며, 중간에 멈춘 학습을 이어 갈 수도 있다.
PyTorch에서 모델 저장과 불러오기의 핵심 도구는 세 가지다. torch.save, torch.load, 그리고 model.load_state_dict다. 공식 문서에서도 PyTorch 모델 저장과 불러오기를 이해할 때 이 세 함수를 핵심으로 소개한다.1
모델을 저장한다는 것은 무엇을 저장한다는 뜻일까
모델을 저장한다고 하면 흔히 "모델 파일 하나를 저장한다"라고 생각한다. 하지만 조금 더 정확히 말하면, 우리는 모델이 학습한 숫자들을 저장한다.
신경망은 층으로 구성되어 있고, 각 층에는 가중치와 편향이 있다. 학습이란 결국 손실을 줄이는 방향으로 이 숫자들을 조금씩 바꾸는 과정이었다. 그러므로 학습이 끝난 모델을 다시 사용하려면, 학습된 가중치와 편향을 파일로 남겨 두면 된다.
PyTorch에서는 이 학습된 상태를 state_dict라는 형태로 다룬다. 이름에서 알 수 있듯이 state_dict는 상태를 담은 딕셔너리다. 더 구체적으로는 각 층의 이름을 키로 삼고, 그 층이 가진 파라미터 텐서를 값으로 담는다. 예를 들어 fc1.weight, fc1.bias, fc2.weight, fc2.bias 같은 이름 아래에 실제 숫자 덩어리들이 들어 있다.
간단한 모델을 하나 만들어 보자.
import torch
import torch.nn as nn
class SimpleNet(nn.Module):
def __init__(self):
super().__init__()
self.net = nn.Sequential(
nn.Linear(4, 16),
nn.ReLU(),
nn.Linear(16, 3)
)
def forward(self, x):
return self.net(x)
model = SimpleNet()
for name, tensor in model.state_dict().items():
print(name, tensor.shape)출력은 대략 다음과 비슷하다.
net.0.weight torch.Size([16, 4])
net.0.bias torch.Size([16])
net.2.weight torch.Size([3, 16])
net.2.bias torch.Size([3])여기서 net.0.weight는 첫 번째 선형층의 가중치다. 입력 4개를 받아 출력 16개를 만들기 때문에 모양이 [16, 4]다. net.2.weight는 두 번째 선형층의 가중치다. 은닉 표현 16개를 받아 클래스 점수 3개를 만들기 때문에 [3, 16] 모양을 가진다.
중요한 점은 state_dict가 모델의 "학습된 숫자들"을 아주 명확한 이름표와 함께 담고 있다는 것이다. 그래서 PyTorch에서는 일반적으로 모델 객체 전체를 저장하기보다 model.state_dict()를 저장하는 방식을 권장한다.1
표준 흐름: state_dict 저장하기
먼저 가장 기본적인 저장 방식을 보자.
PATH = "simple_model.pth"
torch.save(model.state_dict(), PATH)이 한 줄로 모델의 학습된 파라미터가 simple_model.pth 파일에 저장된다. 확장자는 반드시 정해져 있지는 않지만, PyTorch에서는 관례적으로 .pt 또는 .pth를 자주 사용한다.
여기서 저장되는 것은 모델 클래스 코드 전체가 아니다. 저장되는 것은 모델 안에 들어 있던 파라미터와 일부 버퍼다. 버퍼는 학습 가능한 파라미터는 아니지만 모델 상태에 포함되는 값이다. 대표적으로 나중에 배울 배치 정규화의 이동 평균 같은 값이 있다.
이 방식의 장점은 분명하다. 저장 파일이 비교적 명확하고, 모델의 숫자 상태를 다루는 방식이 깔끔하다. 하지만 그 대신 기억해야 할 점도 있다. state_dict만 저장했다면, 불러올 때 모델 구조를 코드로 다시 만들어야 한다.
다시 말해, 저장 파일은 "내용물"이고 모델 클래스는 "그릇"이다. 내용물을 다시 담으려면 같은 모양의 그릇이 필요하다.
state_dict 불러오기
저장한 모델을 다시 불러오는 표준 흐름은 다음과 같다.
model = SimpleNet()
state_dict = torch.load("simple_model.pth", weights_only=True)
model.load_state_dict(state_dict)
model.eval()여기에는 중요한 순서가 있다.
첫째, 먼저 같은 모델 구조를 가진 객체를 만든다.
model = SimpleNet()둘째, 저장된 파일을 메모리로 읽어 온다.
state_dict = torch.load("simple_model.pth", weights_only=True)셋째, 읽어 온 딕셔너리를 모델에 주입한다.
model.load_state_dict(state_dict)여기서 자주 하는 실수가 있다.
# 잘못된 사용
model.load_state_dict("simple_model.pth")load_state_dict는 파일 경로를 받지 않는다. 이미 torch.load로 읽어 온 딕셔너리를 받는다. 즉, torch.load가 파일을 열고, load_state_dict가 그 안의 가중치를 모델에 끼워 넣는다.
공식 튜토리얼에서는 가중치를 불러올 때 weights_only=True를 사용하는 예를 보여 주며, 이는 가중치 로딩에 필요한 동작으로 범위를 제한하는 좋은 습관으로 설명한다.2 이 책에서도 특별한 이유가 없다면 가중치를 불러올 때 weights_only=True를 사용하는 흐름으로 익히자.
마지막 줄의 model.eval()도 중요하다.
model.eval()eval은 evaluation, 즉 평가 모드를 뜻한다. 모델을 추론에 사용할 때는 학습 모드와 다르게 동작해야 하는 층들이 있다. 대표적으로 드롭아웃과 배치 정규화가 그렇다. 드롭아웃은 학습 중 일부 뉴런을 일부러 꺼서 과적합을 줄이는 기법이다. 하지만 추론할 때도 뉴런을 랜덤하게 꺼 버리면 같은 입력에 대해 예측이 흔들릴 수 있다. 배치 정규화도 학습 중에는 현재 배치 통계를 사용하지만, 추론 중에는 학습 과정에서 누적한 통계를 사용해야 한다.
따라서 저장한 모델을 불러와 예측에 사용할 때는 model.eval()을 호출하는 습관을 들여야 한다. PyTorch 문서 역시 추론 전에 model.eval()을 호출하지 않으면 드롭아웃과 배치 정규화 때문에 일관되지 않은 결과가 나올 수 있다고 설명한다.1
추론할 때는 no_grad를 함께 사용한다
불러온 모델로 예측만 할 때는 기울기를 계산할 필요가 없다. 학습할 때는 손실을 줄이기 위해 각 파라미터의 기울기를 계산해야 했지만, 추론에서는 입력을 넣고 결과만 얻으면 된다.
이때 사용하는 것이 torch.no_grad()다.
model.eval()
x = torch.randn(1, 4)
with torch.no_grad():
logits = model(x)
pred = torch.argmax(logits, dim=1)
print(pred)torch.no_grad() 안에서는 PyTorch가 계산 그래프를 만들지 않는다. 계산 그래프는 역전파를 위해 연산 과정을 기록해 둔 구조였다. 추론만 할 때 이 기록은 필요 없으므로 꺼 두는 편이 메모리와 속도 면에서 유리하다.
정리하면, 추론의 기본 자세는 다음과 같다.
model.eval()
with torch.no_grad():
output = model(input)이 두 줄은 앞으로 이미지 분류기, 텍스트 분류기, Transformer 블록을 다룰 때도 계속 등장한다.
모델 전체를 저장하는 방법도 있다
PyTorch에서는 모델 객체 전체를 저장할 수도 있다.
torch.save(model, "simple_model_full.pth")불러올 때는 다음처럼 쓴다.
model = torch.load("simple_model_full.pth", weights_only=False)
model.eval()이 방식은 처음 보면 훨씬 편해 보인다. 모델 구조를 다시 만들지 않아도 되고, load_state_dict를 따로 호출하지 않아도 된다. 모델과 가중치가 한꺼번에 저장되는 느낌이다.
하지만 이 방식에는 조심해야 할 부분이 있다. PyTorch의 전체 모델 저장은 Python의 pickle 직렬화 방식에 의존한다. pickle은 객체 자체를 독립적인 설계도처럼 완전히 저장한다기보다, 불러올 때 필요한 클래스의 위치와 코드 구조에 기대는 방식으로 동작한다. 그래서 모델 클래스를 정의한 파일 위치가 바뀌거나, 클래스 이름이 바뀌거나, 프로젝트 구조가 크게 바뀌면 예전에 저장한 모델을 불러오지 못할 수 있다. PyTorch 문서도 전체 모델 저장 방식이 클래스와 디렉터리 구조에 묶일 수 있다는 단점을 설명한다.1
그래서 이 책에서는 다음 원칙을 기본으로 삼는다.
간단한 실험이나 임시 저장에서는 전체 모델 저장이 편할 수 있다. 하지만 재사용 가능하고 오래 보관할 모델이라면 state_dict를 저장하고, 같은 모델 구조를 코드로 다시 만든 뒤 가중치를 불러오는 방식을 표준으로 익히는 것이 좋다.
체크포인트: 학습을 이어가기 위한 저장
지금까지는 추론에 사용할 모델 가중치를 저장하는 방법을 봤다. 하지만 실제 학습에서는 조금 더 많은 정보를 저장해야 할 때가 많다.
예를 들어 50에폭 동안 학습하기로 했는데, 23에폭이 끝난 뒤 노트북 연결이 끊겼다고 해 보자. 모델 가중치만 저장해 두었다면 23에폭 시점의 모델은 복원할 수 있다. 하지만 학습을 정말로 자연스럽게 이어 가려면 옵티마이저의 상태도 필요하다.
왜 그럴까?
우리가 배운 Adam 같은 최적화 알고리즘은 단순히 현재 기울기만 보고 움직이지 않는다. 이전 기울기의 흐름, 제곱 평균, 학습률과 관련된 내부 상태를 함께 사용한다. 즉, 옵티마이저도 자기만의 기억을 가지고 있다. 모델 가중치만 복원하고 옵티마이저를 새로 만들면, 모델은 23에폭 시점인데 옵티마이저는 0에폭 시점처럼 시작하는 어색한 상황이 된다.
그래서 학습을 재개하려면 보통 다음 정보를 함께 저장한다.
모델의
state_dict옵티마이저의
state_dict현재 에폭 번호
마지막 손실값이나 검증 성능
필요하다면 학습 설정, 클래스 이름, 최고 성능 기록 등 추가 메타데이터
이렇게 여러 정보를 하나의 딕셔너리로 묶어 저장한 파일을 체크포인트라고 부른다. 체크포인트는 게임의 세이브 파일과 비슷하다. 단순히 캐릭터 모습만 저장하는 것이 아니라, 현재 위치, 아이템, 체력, 진행 상황까지 함께 저장해 두는 것이다.
체크포인트 저장하기
체크포인트 저장 코드는 다음처럼 작성할 수 있다.
checkpoint = {
"epoch": epoch,
"model_state_dict": model.state_dict(),
"optimizer_state_dict": optimizer.state_dict(),
"train_loss": train_loss,
"val_loss": val_loss,
}
torch.save(checkpoint, "checkpoint.pth")이 코드는 학습 루프 안에서 보통 한 에폭이 끝날 때 실행한다. 예를 들어 검증 손실이 좋아졌을 때만 저장하도록 만들 수도 있다.
best_val_loss = float("inf")
for epoch in range(num_epochs):
model.train()
for x_batch, y_batch in train_loader:
optimizer.zero_grad()
pred = model(x_batch)
loss = loss_fn(pred, y_batch)
loss.backward()
optimizer.step()
model.eval()
val_loss = 0.0
with torch.no_grad():
for x_batch, y_batch in val_loader:
pred = model(x_batch)
loss = loss_fn(pred, y_batch)
val_loss += loss.item()
val_loss /= len(val_loader)
if val_loss < best_val_loss:
best_val_loss = val_loss
torch.save({
"epoch": epoch,
"model_state_dict": model.state_dict(),
"optimizer_state_dict": optimizer.state_dict(),
"best_val_loss": best_val_loss,
}, "best_checkpoint.pth")
print(f"새로운 최고 모델 저장: epoch={epoch}, val_loss={val_loss:.4f}")이렇게 하면 매 에폭마다 무조건 저장하지 않고, 검증 성능이 좋아졌을 때만 파일을 갱신한다. 실전에서는 매우 흔한 방식이다. 학습이 끝났을 때 마지막 모델이 항상 최고의 모델이라고 보장할 수 없기 때문이다. 어느 순간 이후에는 모델이 훈련 데이터에 지나치게 맞춰져 검증 성능이 나빠질 수 있다. 나중에 과적합을 배울 때 이 현상을 더 자세히 다룰 것이다.
체크포인트 불러오기
체크포인트를 불러올 때도 순서가 중요하다. 먼저 모델과 옵티마이저를 같은 방식으로 만든 다음, 파일에서 상태를 읽어 와 주입한다.
model = SimpleNet()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
checkpoint = torch.load("best_checkpoint.pth", 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_val_loss = checkpoint["best_val_loss"]
print("다음 에폭부터 학습 재개:", start_epoch)추론만 할 것이라면 model.eval()을 호출한다.
model.eval()학습을 이어 갈 것이라면 model.train()을 호출한다.
model.train()이 차이를 가볍게 넘기면 안 된다. 같은 모델이라도 지금 하려는 일이 추론인지, 추가 학습인지에 따라 모드가 달라져야 한다.
공식 문서에서도 일반 체크포인트를 저장할 때 모델의 state_dict뿐 아니라 옵티마이저의 state_dict, 에폭, 손실 같은 정보를 함께 저장해야 한다고 설명한다.1
장치가 달라질 때: CPU와 GPU 사이의 이동
모델 저장과 불러오기에서 자주 만나는 문제가 장치다. 학습은 GPU에서 했는데, 나중에 노트북 CPU에서 불러와 추론하고 싶을 수 있다. 또는 CPU에서 저장한 작은 모델을 GPU 서버로 옮겨 추가 학습하고 싶을 수도 있다.
PyTorch 텐서는 어느 장치에 올라가 있는지 정보를 가지고 있다. 그래서 저장 당시의 장치와 불러오는 장치가 다르면 조심해야 한다. 이때 사용하는 인자가 map_location이다.
GPU에서 저장한 모델을 CPU에서 불러오려면 다음처럼 쓴다.
device = torch.device("cpu")
model = SimpleNet()
state_dict = torch.load(
"simple_model.pth",
map_location=device,
weights_only=True
)
model.load_state_dict(state_dict)
model.to(device)
model.eval()반대로 GPU가 가능하면 GPU로 불러와 사용할 수 있다.
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = SimpleNet()
state_dict = torch.load(
"simple_model.pth",
map_location=device,
weights_only=True
)
model.load_state_dict(state_dict)
model.to(device)
model.eval()여기서 model.to(device)만으로 끝나지 않는다. 입력 텐서도 같은 장치로 옮겨야 한다.
x = torch.randn(1, 4).to(device)
with torch.no_grad():
output = model(x)모델은 GPU에 있는데 입력은 CPU에 있거나, 반대로 모델은 CPU에 있는데 입력은 GPU에 있으면 오류가 발생한다. PyTorch는 "서로 다른 작업대에 놓인 재료를 한 번에 계산할 수 없다"라고 말하는 셈이다. 모델과 입력은 같은 장치에 있어야 한다.
작은 실습: 저장한 모델로 다시 예측하기
이제 전체 흐름을 하나의 짧은 코드로 묶어 보자. 여기서는 이미 학습이 끝난 model이 있다고 가정한다.
import torch
import torch.nn as nn
class SimpleNet(nn.Module):
def __init__(self):
super().__init__()
self.net = nn.Sequential(
nn.Linear(4, 16),
nn.ReLU(),
nn.Linear(16, 3)
)
def forward(self, x):
return self.net(x)
# 1. 학습된 모델이 있다고 가정
model = SimpleNet()
# 실제 상황에서는 여기에서 학습 루프가 실행되었다고 생각하자.
# 2. 가중치 저장
torch.save(model.state_dict(), "simple_model.pth")
# 3. 새 모델 객체 생성
loaded_model = SimpleNet()
# 4. 저장된 가중치 불러오기
state_dict = torch.load("simple_model.pth", weights_only=True)
loaded_model.load_state_dict(state_dict)
# 5. 추론 모드 전환
loaded_model.eval()
# 6. 예측
sample = torch.randn(1, 4)
with torch.no_grad():
logits = loaded_model(sample)
prediction = torch.argmax(logits, dim=1)
print("예측 클래스:", prediction.item())이 코드는 작지만 중요한 구조를 모두 담고 있다. 저장하고, 같은 모델 구조를 만들고, 가중치를 주입하고, 평가 모드로 바꾼 뒤, 기울기 계산 없이 추론한다.
앞으로 우리가 만들 숫자 분류기, 이미지 분류기, 텍스트 분류기도 이 흐름을 따른다. 모델이 커지고 데이터가 복잡해져도 기본 원리는 변하지 않는다.
저장 파일을 다룰 때의 실전 감각
모델 저장을 처음 배울 때는 코드 한두 줄만 외우기 쉽다. 하지만 실제 프로젝트에서는 저장 파일 관리도 중요하다. 파일 이름을 아무렇게나 만들면 나중에 어떤 모델이 어떤 실험에서 나온 것인지 알 수 없다.
좋은 습관은 저장 파일 이름이나 체크포인트 딕셔너리에 맥락을 남기는 것이다. 예를 들어 다음처럼 저장할 수 있다.
torch.save({
"epoch": epoch,
"model_state_dict": model.state_dict(),
"optimizer_state_dict": optimizer.state_dict(),
"best_val_loss": best_val_loss,
"learning_rate": 0.001,
"model_name": "SimpleNet",
}, "simplenet_best.pth")이렇게 해 두면 몇 주 뒤에 파일을 열어도 어느 정도 상황을 복원할 수 있다. 딥러닝 실험은 생각보다 기억에 의존하기 어렵다. 오늘은 너무나 당연해 보이는 설정도, 다음 주에는 희미해진다.
특히 다음 세 가지는 자주 기록해 두는 편이 좋다.
모델 구조를 구분할 수 있는 이름
학습률, 배치 크기, 에폭 수 같은 주요 학습 설정
검증 손실이나 정확도처럼 모델 선택에 사용한 성능 지표
모델 저장은 단순히 파일을 만드는 일이 아니라 실험의 흔적을 남기는 일이다.
자주 만나는 오류들
모델을 저장하고 불러오는 과정에서 가장 자주 만나는 오류는 구조 불일치다. state_dict에는 fc1.weight라는 이름의 텐서가 들어 있는데, 새로 만든 모델에는 linear1.weight라는 이름만 있다면 PyTorch는 어디에 이 가중치를 넣어야 할지 모른다.
또는 이름은 같지만 텐서 모양이 다를 수도 있다. 예전 모델의 마지막 출력이 10개 클래스였는데, 새 모델은 3개 클래스로 바뀌었다면 마지막 층의 가중치 모양이 맞지 않는다.
이럴 때는 저장 당시의 모델 클래스와 현재 모델 클래스가 같은지 확인해야 한다. state_dict 방식은 유연하지만, 마법은 아니다. 같은 구조의 그릇이 있어야 같은 내용물을 담을 수 있다.
다음으로 자주 만나는 문제는 eval()을 빼먹는 것이다. 특히 드롭아웃이 들어간 모델에서 이 실수는 예측 결과를 흔들리게 만든다. "왜 같은 이미지를 넣었는데 결과가 조금씩 달라지지?"라는 질문이 생기면 먼저 model.eval()을 확인하자.
마지막으로 장치 불일치 오류가 있다. 오류 메시지에 cpu와 cuda가 함께 보인다면 모델과 입력 텐서가 같은 장치에 있는지 확인해야 한다.
모델 저장은 딥러닝 작업의 마침표이자 쉼표다
모델을 저장하는 일은 학습의 마침표처럼 보인다. "여기까지 학습했다"라고 결과를 고정하기 때문이다. 하지만 동시에 쉼표이기도 하다. 저장한 모델을 나중에 다시 불러와 추론할 수 있고, 더 많은 데이터로 추가 학습할 수 있으며, 다른 모델의 출발점으로 사용할 수도 있다.
이 장에서 배운 흐름을 다시 연결해 보자.
모델이 예측한다. 손실을 계산한다. 역전파로 각 가중치가 오차에 얼마나 영향을 줬는지 본다. 옵티마이저가 가중치를 조금씩 수정한다. 이 과정을 반복한다. 그리고 어느 순간, 그 결과를 state_dict라는 형태로 붙잡아 파일에 저장한다.
이제 우리는 학습 루프를 만들 수 있을 뿐 아니라, 그 결과를 보관하고 되살릴 수 있다. 다음 장부터는 이 능력을 바탕으로 본격적인 모델을 만들기 시작한다. 먼저 가장 고전적이고도 중요한 실습인 손글씨 숫자 분류로 들어갈 것이다. 숫자 이미지를 보고 0부터 9까지의 클래스를 맞히는 모델을 만들면서, 지금까지 배운 텐서, 데이터로더, 학습 루프, 모델 저장이 하나의 완성된 작업으로 이어지는 모습을 보게 될 것이다.
참고
이 노트는 보도·비평·교육·연구 목적으로 작성되었습니다. 저작권 문의가 있으시면 에서 알려주세요.