좋은 PyTorch 코드 습관으로 마무리하기

작동하는 코드에서 재현 가능한 코드로

PyTorch를 처음 배울 때는 "일단 실행되는 코드"가 중요합니다. 하지만 다음 단계로 가려면 같은 코드를 다시 실행했을 때 비슷한 결과가 나오고, 나중에 읽어도 의도를 알 수 있어야 합니다. 공식 튜토리얼도 데이터 준비, 모델 생성, 최적화, 저장으로 이어지는 전체 흐름을 반복해서 강조합니다.¹

가장 먼저 실행 환경을 코드에 남기세요. CPU, CUDA, MPS 중 어떤 장치를 쓸지 정하고, 난수 시드와 하이퍼파라미터를 한곳에 모읍니다.

import random
import numpy as np
import torch

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

CONFIG = {
    "seed": 42,
    "batch_size": 64,
    "lr": 1e-3,
    "epochs": 10,
    "model_path": "checkpoints/best_model.pth",
}

def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)

여기에 PyTorch 버전, 데이터 분할 기준, 사용한 모델 파일 이름도 함께 기록하면 좋습니다.

print(torch.__version__)
print(DEVICE)

다만 시드를 고정해도 모든 환경에서 결과가 완전히 똑같아지는 것은 아닙니다. GPU 연산, 라이브러리 버전, 일부 비결정적 연산 때문에 작은 차이가 날 수 있습니다. 그래도 시드와 설정을 기록해 두면 실험이 실패했을 때 "무엇이 달라졌는지" 찾기가 훨씬 쉬워집니다.

파일을 역할별로 나누기

코드가 길어지면 한 파일에 모두 넣지 마세요. 작은 프로젝트라도 다음처럼 나누면 읽기 쉬워집니다.

project/
├── model.py      # nn.Module 모델 정의
├── dataset.py    # Dataset, Transform, DataLoader
├── train.py      # 학습과 평가 실행
├── utils.py      # seed, accuracy, save/load 함수
└── checkpoints/  # 저장된 모델

모델은 nn.Module로 만들고, 데이터 처리는 Dataset과 DataLoader에 맡기고, 학습 루프는 함수로 분리합니다. 이렇게 하면 모델만 바꾸거나 데이터셋만 바꾸는 실험이 쉬워집니다. 노트북은 아이디어를 빠르게 확인할 때 좋고, 오래 쓸 코드는 .py 파일로 옮기는 습관이 좋습니다.

역할을 나누는 기준은 "나중에 바꾸고 싶은 부분이 어디인가?"입니다. 모델 구조를 자주 바꿀 것 같다면 model.py로 분리하고, 데이터 전처리를 자주 실험할 것 같다면 dataset.py를 따로 두는 식입니다. 처음부터 복잡한 구조를 만들 필요는 없지만, 코드가 길어질수록 책임을 나누는 습관이 중요해집니다.

학습과 평가를 분명히 구분하기

학습할 때는 model.train()을, 검증과 추론에는 model.eval()을 사용합니다. 평가 중에는 기울기가 필요 없으므로 torch.no_grad()나 torch.inference_mode()를 함께 씁니다.

def evaluate(model, dataloader, loss_fn):
    model.eval()
    total_loss, correct, total = 0, 0, 0

    with torch.inference_mode():
        for x, y in dataloader:
            x, y = x.to(DEVICE), y.to(DEVICE)
            pred = model(x)
            loss = loss_fn(pred, y)

            total_loss += loss.item() * x.size(0)
            correct += (pred.argmax(1) == y).sum().item()
            total += y.size(0)

    return total_loss / total, correct / total

torch.inference_mode()는 추론 전용 상황에서 torch.no_grad()보다 더 강하게 autograd 관련 처리를 끄는 방식입니다. 평가와 추론 코드에서는 좋은 선택입니다. 반대로, 나중에 다시 기울기 계산에 연결해야 하는 텐서를 만들고 있다면 no_grad()와의 차이를 확인해야 합니다. 초보 단계에서는 "일반 평가와 추론에는 inference_mode() 또는 no_grad()를 사용한다" 정도로 기억하면 충분합니다.

체크포인트도 습관화하세요. 예측만 할 모델이라면 model.state_dict()를 저장하면 됩니다. 학습을 이어서 실행해야 한다면 모델 가중치뿐 아니라 옵티마이저 상태, 에폭, 최고 정확도도 함께 저장합니다.

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

결과를 눈으로 확인하기

손실과 정확도는 숫자로만 보지 마세요. tqdm으로 진행 상황을 확인하고, TensorBoard나 matplotlib으로 학습 곡선을 그려 보세요. PyTorch 한국어 튜토리얼에도 TensorBoard로 모델, 데이터, 학습을 시각화하는 내용이 포함되어 있습니다.²

학습 손실은 계속 내려가는데 검증 손실이 올라가면 과적합일 수 있습니다. 두 손실이 모두 줄지 않으면 학습률, 데이터 전처리, 출력 차원, 라벨 타입을 의심해야 합니다. 샘플 이미지를 몇 개 직접 예측해 보는 것도 좋은 디버깅 방법입니다.

실험을 비교할 때는 기록 형식을 정해 두면 좋습니다.

실험 이름: cnn_lr_001
모델: SmallCNN
데이터: FashionMNIST train/val split seed=42
batch_size: 64
learning_rate: 0.001
best_val_acc: 0.XXXX
메모: 5 epoch 이후 검증 정확도 정체

처음에는 이런 기록이 번거롭게 느껴질 수 있습니다. 하지만 실험이 세 개만 넘어가도 "어떤 설정이 가장 좋았는지" 기억에만 의존하기 어렵습니다.

예제 코드는 검증해서 가져오기

새 모델이나 데이터셋을 만났을 때는 공식 문서, 한국어 튜토리얼, TorchVision 예제, 신뢰할 수 있는 공개 저장소를 함께 확인하세요. 단, 예제 코드를 그대로 붙여 넣으면 자주 실패합니다. 입력 이미지 크기, 채널 수, 정규화 방식, 클래스 개수, 출력 차원이 내 문제와 맞는지 반드시 확인해야 합니다.

가져온 코드를 실행하기 전에는 다음을 먼저 확인하세요.

• 모델 첫 계층의 입력 채널 수가 내 데이터와 맞는가?

• 마지막 출력 차원이 클래스 개수와 같은가?

• 예제의 정규화 평균과 표준편차를 내 데이터에 그대로 써도 되는가?

• 손실 함수가 요구하는 라벨 형식과 내 라벨이 맞는가?

• 저장된 가중치를 불러올 때 모델 구조가 같은가?

초보 이후에는 전이학습, TorchVision 사전학습 모델, 혼합 정밀도 학습, 효율적인 DataLoader 설정, torch.compile, torch.export, ONNX 내보내기 같은 주제로 넓혀 가면 좋습니다. 특히 최신 PyTorch에서는 배포와 내보내기 흐름에서 torch.export의 중요성이 커지고 있으며, TorchScript는 새 프로젝트의 기본 선택으로 권장하기 어려운 방향으로 바뀌고 있습니다.³

마지막으로 기억할 원칙은 단순합니다. 좋은 PyTorch 코드는 빠르게만 도는 코드가 아닙니다. 다시 실행할 수 있고, 고치기 쉽고, 결과를 설명할 수 있는 코드입니다. 이 습관이 쌓이면 새 모델을 만났을 때도 흔들리지 않습니다.

참고

¹Learn the Basics — PyTorch Tutorials

²파이토치 한국어 튜토리얼

³PyTorch 2.12 Release Blog — PyTorch