PyTorch 입문(Introduction to PyTorch)

피스톤과 크랭크축으로 엔진을 직접 만들었다면, 이제 대부분의 사람이 실제로 운전하는 차를 익힙니다.

유형: Build 언어: Python 선수 학습: Lesson 03.10 Build Your Own Mini Framework 소요 시간: 약 75분

학습 목표

• PyTorch의 nn.Module, nn.Sequential, 오토그래드(Autograd)를 사용해 신경망(Neural Network)을 만들고 학습합니다.
• PyTorch 텐서(Tensor), GPU 가속, 표준 학습 루프(zero_grad, forward, loss, backward, step)를 사용합니다.
• 직접 만든 미니 프레임워크(mini framework)의 구성 요소를 PyTorch의 대응 개념으로 옮겨 봅니다.
• 같은 문제에서 순수 Python 프레임워크와 PyTorch의 학습 속도를 프로파일링(Profiling)하고 비교합니다.

문제

이전 lesson에서 동작하는 미니 프레임워크(mini framework)를 만들었습니다. 선형 레이어(Linear layer), ReLU, 드롭아웃(Dropout), 배치 정규화(Batch Normalization), Adam, DataLoader, 학습 루프까지 있습니다. 순수 Python만으로 원형 분류 문제에 4-레이어 신경망(4-layer network)을 학습시킬 수 있습니다.

하지만 같은 문제에서 PyTorch보다 약 500배 느립니다.

직접 만든 미니 프레임워크는 중첩된 Python 반복문(loop)으로 샘플을 하나씩 처리합니다. PyTorch는 같은 연산을 최적화된 C++/CUDA 커널(Kernel)로 보냅니다. GPU에서는 이 커널들이 병렬로 실행됩니다. 단일 NVIDIA A100에서 PyTorch는 ImageNet의 128만 장 이미지로 ResNet-50(2,560만 파라미터)을 약 6시간 안에 학습할 수 있습니다. 같은 일을 직접 만든 프레임워크로 한다면 메모리가 먼저 부족하지 않더라도 약 3,000시간이 걸릴 수 있습니다.

속도만 차이가 아닙니다. 직접 만든 프레임워크에는 GPU 지원이 없습니다. 모든 모듈(Module)마다 backward()를 손으로 구현해야 하므로 자동 미분(Automatic Differentiation)도 없습니다. 직렬화(Serialization), 분산 학습(Distributed Training), 혼합 정밀도(Mixed Precision)도 없고, print 문 없이 기울기 흐름(gradient flow)을 디버깅할 방법도 없습니다.

PyTorch는 이 차이를 메워 줍니다. 동시에 이미 직접 만든 사고 모델을 그대로 유지합니다. 모듈(Module), forward(), parameters(), backward(), optimizer.step()이라는 개념은 거의 일대일로 이어집니다. 문법도 비슷합니다. 차이는 PyTorch가 같은 인터페이스(interface) 뒤에 수년간의 시스템 엔지니어링을 감춰 둔다는 점입니다.

개념

PyTorch가 널리 쓰이게 된 이유

2015년의 TensorFlow는 무엇을 실행하기 전에 정적 계산 그래프(Static Computation Graph)를 먼저 정의해야 했습니다. 그래프를 만들고, 컴파일한 뒤, 데이터를 흘려보내야 했습니다. 디버깅은 그래프 시각화를 들여다보는 일이었고, 아키텍처(architecture)를 바꾸려면 그래프를 다시 만들어야 했습니다.

PyTorch는 2017년에 즉시 실행(Eager Execution)이라는 다른 철학으로 등장했습니다. Python 코드를 쓰면 바로 실행됩니다. y = model(x)는 "나중에 y를 계산할 노드를 그래프에 추가"하는 것이 아니라 지금 실제로 y를 계산합니다. 그래서 표준 Python 디버깅 도구가 그대로 동작했습니다. print(), pdb, breakpoint(), 그리고 순전파(forward pass) 안의 if/else가 모두 자연스럽게 작동했습니다.

2020년 무렵에는 연구 생태계가 PyTorch 쪽으로 크게 이동했습니다. 머신러닝(ML) 논문에서 PyTorch 점유율은 2017년 7% 수준에서 2022년 75% 이상으로 커졌습니다. Meta, Google DeepMind, OpenAI, Anthropic, Hugging Face가 모두 PyTorch를 핵심 프레임워크(framework)로 사용합니다. TensorFlow 2.x도 이후 즉시 실행(eager execution)을 도입했습니다.

교훈은 명확합니다. 개발자 경험은 복리로 쌓입니다. 10% 느리더라도 50% 더 빨리 디버깅할 수 있는 프레임워크는 결국 이깁니다.

텐서(Tensor)

텐서(Tensor)는 세 가지 중요한 속성을 가진 다차원 배열입니다. 모양(Shape), 자료형(Dtype), 장치(Device)입니다.

import torch

x = torch.zeros(3, 4)           # 모양: (3, 4), 자료형: float32, 장치: cpu
x = torch.randn(2, 3, 224, 224) # 224x224 RGB 이미지 2장의 배치
x = torch.tensor([1, 2, 3])     # Python 리스트에서 생성

모양(Shape)은 차원을 뜻합니다. 스칼라(Scalar)는 (), 벡터(Vector)는 (n,), 행렬(Matrix)은 (m, n), 이미지 배치는 (batch, channels, height, width)입니다.

자료형(Dtype)은 정밀도와 메모리 사용량을 결정합니다.

장치(Device)는 연산이 어디에서 일어나는지 결정합니다.

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
x = torch.randn(3, 4, device=device)
x = x.to("cuda")
x = x.cpu()

모든 연산은 관련 텐서가 같은 장치에 있어야 합니다. PyTorch 초보자가 가장 자주 만나는 오류는 RuntimeError: Expected all tensors to be on the same device입니다. 계산 전에 모델, 입력, 라벨을 같은 장치(device)로 옮겨야 합니다.

모양 변경(Reshaping)은 보통 데이터 자체를 복사하지 않고 메타데이터(metadata)만 바꾸므로 상수 시간에 가깝습니다.

x = torch.randn(2, 3, 4)
x.view(2, 12)      # (2, 12)로 변경, 메모리 연속(contiguous) 조건 필요
x.reshape(6, 4)    # (6, 4)로 변경, 대부분의 경우 동작
x.permute(2, 0, 1) # 차원 순서 변경
x.unsqueeze(0)     # 차원 추가: (1, 2, 3, 4)
x.squeeze()        # 크기가 1인 차원 제거

오토그래드(Autograd)

직접 만든 미니 프레임워크에서는 모듈마다 backward()를 구현해야 했습니다. PyTorch는 그럴 필요가 없습니다. PyTorch는 텐서에 수행된 모든 연산을 방향 비순환 그래프(Directed Acyclic Graph; DAG), 즉 계산 그래프(Computational Graph)에 기록한 뒤, 그 그래프를 역방향으로 따라가며 기울기(gradient)를 자동 계산합니다.

graph LR
    x["x (leaf)"] --> mul["*"]
    w["w (leaf, requires_grad)"] --> mul
    mul --> add["+"]
    b["b (leaf, requires_grad)"] --> add
    add --> loss["loss"]
    loss --> |".backward()"| add
    add --> |"grad"| b
    add --> |"grad"| mul
    mul --> |"grad"| w

핵심 차이는 PyTorch가 테이프 기반 자동 미분(Tape-based Autodiff)을 쓴다는 점입니다. 순전파(forward pass) 중 모든 연산이 테이프(tape)에 기록됩니다. .backward()를 호출하면 이 테이프를 역순으로 재생하면서 기울기를 계산합니다.

x = torch.randn(3, requires_grad=True)
y = x ** 2 + 3 * x
z = y.sum()
z.backward()
print(x.grad)  # dz/dx = 2x + 3

오토그래드의 세 가지 규칙은 다음과 같습니다.

1. requires_grad=True인 잎 텐서(leaf tensor)만 기울기를 누적합니다.
2. 기울기는 기본적으로 누적됩니다. 매 역전파(backward pass) 전에 optimizer.zero_grad()를 호출합니다.
3. torch.no_grad()는 기울기 추적(gradient tracking)을 비활성화합니다. 평가(Evaluation)와 추론(Inference)에 사용합니다.

nn.Module

nn.Module은 PyTorch의 모든 신경망 구성 요소가 상속하는 기본 클래스(base class)입니다. Lesson 10에서 이미 이 추상화(abstraction)를 직접 만들었습니다. PyTorch의 버전은 자동 파라미터(parameter) 등록, 재귀적인 모듈(module) 탐색, 장치(device) 관리, state_dict 직렬화(serialization)를 추가로 제공합니다.

import torch.nn as nn

class MLP(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super().__init__()
        self.layer1 = nn.Linear(input_dim, hidden_dim)
        self.relu = nn.ReLU()
        self.layer2 = nn.Linear(hidden_dim, output_dim)

    def forward(self, x):
        x = self.layer1(x)
        x = self.relu(x)
        x = self.layer2(x)
        return x

__init__에서 nn.Module이나 nn.Parameter를 속성으로 할당하면 PyTorch가 자동으로 등록합니다. model.parameters()는 등록된 모든 파라미터를 재귀적으로 모읍니다. 그래서 미니 프레임워크처럼 가중치(weight)를 직접 모을 필요가 없습니다.

손실 함수(Loss Functions)와 옵티마이저(Optimizers)

PyTorch는 앞 lesson에서 직접 만든 것들의 운영 가능한(production-ready) 버전을 제공합니다.

손실 함수(Loss function)는 torch.nn에 있습니다.

CrossEntropyLoss는 내부적으로 LogSoftmax와 NLLLoss를 결합합니다. 따라서 소프트맥스 결과가 아니라 가공 전 로짓(raw logits)을 넣어야 합니다. 이 실수는 잘못된 기울기를 조용히 만들 수 있습니다.

옵티마이저(Optimizer)는 torch.optim에 있습니다.

학습 루프(Training Loop)

PyTorch 학습 루프는 거의 항상 같은 5단계 패턴을 따릅니다. Lesson 10에서 이미 본 흐름입니다.

sequenceDiagram
    participant D as DataLoader
    participant M as Model
    participant L as Loss fn
    participant O as Optimizer

    loop Each Epoch
        D->>M: batch = next(dataloader)
        M->>L: predictions = model(batch)
        L->>L: loss = criterion(predictions, targets)
        L->>M: loss.backward()
        O->>M: optimizer.step()
        O->>O: optimizer.zero_grad()
    end

표준 패턴은 다음과 같습니다.

for epoch in range(num_epochs):
    model.train()
    for inputs, targets in train_loader:
        inputs, targets = inputs.to(device), targets.to(device)
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()

배치 반복문(batch loop) 안의 다섯 줄입니다. GPT-4, Stable Diffusion, LLaMA도 결국 이 패턴의 변형으로 학습됩니다. 아키텍처와 데이터는 바뀌지만 이 골격은 크게 바뀌지 않습니다.

데이터셋(Dataset)과 데이터로더(DataLoader)

PyTorch의 Dataset은 __len__, __getitem__ 두 메서드를 가진 추상 클래스입니다. DataLoader는 여기에 배치 구성(batching), 섞기(shuffling), 다중 프로세스 데이터 적재(multi-process data loading)를 얹습니다.

from torch.utils.data import Dataset, DataLoader

class MNISTDataset(Dataset):
    def __init__(self, images, labels):
        self.images = images
        self.labels = labels

    def __len__(self):
        return len(self.labels)

    def __getitem__(self, idx):
        return self.images[idx], self.labels[idx]

loader = DataLoader(dataset, batch_size=64, shuffle=True, num_workers=4)

num_workers=4는 4개의 프로세스로 데이터를 병렬 적재합니다. GPU가 현재 배치(batch)를 학습하는 동안 다음 배치를 준비할 수 있습니다. 큰 이미지나 오디오처럼 디스크 입출력에 묶인(disk-bound) 작업에서는 이것만으로도 학습 속도가 크게 빨라질 수 있습니다.

GPU 학습

모델을 GPU로 옮기는 방법은 다음과 같습니다.

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

이 코드는 모든 파라미터와 버퍼(buffer)를 재귀적으로 GPU로 옮깁니다. 학습 중에는 배치도 같은 장치로 옮깁니다.

inputs, targets = inputs.to(device), targets.to(device)

혼합 정밀도(Mixed Precision)는 최신 GPU(A100, H100, RTX 4090 등)에서 순전파/역전파(forward/backward)를 float16으로 실행하고 마스터 가중치(master weight)는 float32로 유지해 메모리 사용량을 줄이고 처리량(throughput)을 높입니다.

from torch.amp import autocast, GradScaler

scaler = GradScaler()
for inputs, targets in loader:
    with autocast(device_type="cuda"):
        outputs = model(inputs)
        loss = criterion(outputs, targets)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()
    optimizer.zero_grad()

비교: 미니 프레임워크(Mini Framework) vs PyTorch vs JAX

만들어 보기

PyTorch의 기본 요소(primitive)만 사용해 MNIST에 3-레이어 MLP(3-layer MLP)를 학습시킵니다. 상위 수준 래퍼(high-level wrapper)도, torchvision.datasets도 쓰지 않습니다. 원본 데이터(raw data)를 직접 다운로드하고 파싱(parsing)합니다.

1단계: 원본 파일에서 MNIST 불러오기

MNIST는 gzip으로 압축된 4개 파일로 제공됩니다. 학습용 이미지(60,000 x 28 x 28), 학습용 라벨, 테스트용 이미지(10,000 x 28 x 28), 테스트용 라벨입니다. 이를 다운로드하고 이진 형식(binary format)을 파싱합니다.

import torch
import torch.nn as nn
import struct
import gzip
import urllib.request
import os

def download_mnist(path="./mnist_data"):
    base_url = "https://storage.googleapis.com/cvdf-datasets/mnist/"
    files = [
        "train-images-idx3-ubyte.gz",
        "train-labels-idx1-ubyte.gz",
        "t10k-images-idx3-ubyte.gz",
        "t10k-labels-idx1-ubyte.gz",
    ]
    os.makedirs(path, exist_ok=True)
    for f in files:
        filepath = os.path.join(path, f)
        if not os.path.exists(filepath):
            urllib.request.urlretrieve(base_url + f, filepath)

def load_images(filepath):
    with gzip.open(filepath, "rb") as f:
        magic, num, rows, cols = struct.unpack(">IIII", f.read(16))
        data = f.read()
        images = torch.frombuffer(bytearray(data), dtype=torch.uint8)
        images = images.reshape(num, rows * cols).float() / 255.0
    return images

def load_labels(filepath):
    with gzip.open(filepath, "rb") as f:
        magic, num = struct.unpack(">II", f.read(8))
        data = f.read()
        labels = torch.frombuffer(bytearray(data), dtype=torch.uint8).long()
    return labels

2단계: 모델 정의하기

3-레이어 MLP입니다. 구조는 784 -> 256 -> 128 -> 10입니다. ReLU 활성화 함수(activation)와 정칙화(regularization)를 위한 드롭아웃을 사용합니다. 단순하게 유지하기 위해 배치 정규화(batch norm)는 넣지 않습니다.

class MNISTModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(784, 256),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(128, 10),
        )

    def forward(self, x):
        return self.net(x)

출력 레이어(output layer)는 숫자 0~9에 해당하는 10개의 가공 전 로짓(raw logits)을 만듭니다. 소프트맥스는 쓰지 않습니다. CrossEntropyLoss가 내부에서 처리합니다.

파라미터 수는 784*256 + 256 + 256*128 + 128 + 128*10 + 10 = 235,146입니다. 현대 기준으로는 매우 작습니다. GPT-2 small은 1억 2,400만(124M) 파라미터입니다. 이 모델은 몇 초 안에 학습됩니다.

3단계: 학습 루프

표준 순전파-손실-역전파-업데이트(forward-loss-backward-step) 패턴입니다.

def train_one_epoch(model, loader, criterion, optimizer, device):
    model.train()
    total_loss = 0
    correct = 0
    total = 0
    for images, labels in loader:
        images, labels = images.to(device), labels.to(device)
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * images.size(0)
        _, predicted = outputs.max(1)
        correct += predicted.eq(labels).sum().item()
        total += labels.size(0)
    return total_loss / total, correct / total

평가에서는 torch.no_grad()를 사용합니다. 이것은 오토그래드를 비활성화해 메모리를 줄이고 추론(inference)을 빠르게 합니다. 비활성화하지 않으면 PyTorch는 사용하지 않는 계산 그래프를 계속 만들어 둡니다.

def evaluate(model, loader, criterion, device):
    model.eval()
    total_loss = 0
    correct = 0
    total = 0
    with torch.no_grad():
        for images, labels in loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            loss = criterion(outputs, labels)
            total_loss += loss.item() * images.size(0)
            _, predicted = outputs.max(1)
            correct += predicted.eq(labels).sum().item()
            total += labels.size(0)
    return total_loss / total, correct / total

4단계: 전체 연결하기

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

    download_mnist()
    train_images = load_images("./mnist_data/train-images-idx3-ubyte.gz")
    train_labels = load_labels("./mnist_data/train-labels-idx1-ubyte.gz")
    test_images = load_images("./mnist_data/t10k-images-idx3-ubyte.gz")
    test_labels = load_labels("./mnist_data/t10k-labels-idx1-ubyte.gz")

    train_dataset = torch.utils.data.TensorDataset(train_images, train_labels)
    test_dataset = torch.utils.data.TensorDataset(test_images, test_labels)
    train_loader = torch.utils.data.DataLoader(
        train_dataset, batch_size=64, shuffle=True
    )
    test_loader = torch.utils.data.DataLoader(
        test_dataset, batch_size=256, shuffle=False
    )

    model = MNISTModel().to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

    num_params = sum(p.numel() for p in model.parameters())
    print(f"Device: {device}")
    print(f"Parameters: {num_params:,}")
    print(f"Train samples: {len(train_dataset):,}")
    print(f"Test samples: {len(test_dataset):,}")
    print()

    for epoch in range(10):
        train_loss, train_acc = train_one_epoch(
            model, train_loader, criterion, optimizer, device
        )
        test_loss, test_acc = evaluate(
            model, test_loader, criterion, device
        )
        print(
            f"Epoch {epoch+1:2d} | "
            f"Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.4f} | "
            f"Test Loss: {test_loss:.4f} | Test Acc: {test_acc:.4f}"
        )

    torch.save(model.state_dict(), "mnist_mlp.pt")
    print(f"\nModel saved to mnist_mlp.pt")
    print(f"Final test accuracy: {test_acc:.4f}")

10 에포크(epoch) 후 기대되는 테스트 정확도(test accuracy)는 약 97.8%입니다. CPU에서는 약 30초, GPU에서는 약 5초가 걸립니다. 같은 아키텍처를 직접 만든 미니 프레임워크로 실행하면 약 45분이 걸릴 수 있습니다.

사용하기

빠른 비교: 미니 프레임워크(Mini Framework) vs PyTorch

인터페이스(interface)는 거의 같습니다. 차이는 내부 구현입니다.

모델 저장과 로딩

torch.save(model.state_dict(), "model.pt")

model = MNISTModel()
model.load_state_dict(torch.load("model.pt", weights_only=True))
model.eval()

항상 모델 객체가 아니라 state_dict()(파라미터 딕셔너리)를 저장합니다. 모델 객체 저장은 pickle을 사용하므로 코드를 리팩터링(refactoring)하면 깨질 수 있습니다. state_dict는 더 이식 가능하고 안정적인 표현입니다.

학습률 스케줄링(Learning Rate Scheduling)

scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
    optimizer, T_max=10
)
for epoch in range(10):
    train_one_epoch(model, train_loader, criterion, optimizer, device)
    scheduler.step()

PyTorch는 StepLR, ExponentialLR, CosineAnnealingLR, OneCycleLR, ReduceLROnPlateau 등 15개 이상의 스케줄러(scheduler)를 제공합니다. 모두 같은 옵티마이저 인터페이스에 연결됩니다.

산출물 만들기

이 lesson의 최종 산출물은 다음 두 가지입니다.

• outputs/prompt-pytorch-debugger.md: PyTorch 학습 실패를 진단하는 prompt
• outputs/skill-pytorch-patterns.md: PyTorch 학습 패턴을 정리한 skill 참고 자료

연습문제

1. (쉬움) 배치 정규화(Batch normalization)를 추가합니다. 각 선형 레이어(linear layer) 뒤, 활성화 함수(activation) 앞에 nn.BatchNorm1d를 넣습니다. 드롭아웃만 사용한 버전과 테스트 정확도(test accuracy), 학습 속도(training speed)를 비교합니다. 배치 정규화는 더 적은 에포크 안에 98% 이상에 도달할 수 있습니다.
2. (중간) 학습률 탐색기(learning rate finder)를 구현합니다. 한 에포크 동안 학습률을 1e-7에서 1.0까지 지수적으로 증가시키며 학습합니다. 손실(loss) 대 학습률(LR) 그래프를 그립니다. 손실이 올라가기 직전의 값을 좋은 학습률 후보로 사용합니다. 이를 활용해 MNIST 모델에 더 나은 학습률을 정합니다.
3. (중간) GPU 혼합 정밀도(mixed precision)로 옮깁니다. 학습 루프에 torch.amp.autocast와 GradScaler를 추가합니다. GPU에서 혼합 정밀도 적용 전후의 처리량(samples/second)을 측정합니다. A100에서는 약 2배 속도 향상을 기대할 수 있습니다.
4. (중간) 사용자 정의(custom) Dataset을 만듭니다. MNIST와 같은 형식이지만 의류 이미지가 들어 있는 Fashion-MNIST를 다운로드합니다. FashionMNISTDataset(Dataset)을 __getitem__, __len__로 구현하고 같은 MLP를 학습합니다. Fashion-MNIST는 더 어렵기 때문에 약 88% 정확도를 기대합니다.
5. (어려움) Adam을 SGD + 모멘텀(momentum)으로 교체합니다. SGD(params, lr=0.01, momentum=0.9)로 학습합니다. 수렴 곡선(convergence curve)을 비교합니다. 그 다음 CosineAnnealingLR 스케줄러를 추가해 10 에포크 안에 SGD가 Adam을 따라잡는지 확인합니다.

핵심 용어

더 읽을거리

• Paszke et al., "PyTorch: An Imperative Style, High-Performance Deep Learning Library" (2019) — PyTorch의 설계 절충점을 설명한 원논문
• PyTorch Tutorials: "Learning PyTorch with Examples" — 텐서에서 nn.Module까지 이어지는 공식 입문 경로
• PyTorch Performance Tuning Guide — 혼합 정밀도, DataLoader 워커, 고정 메모리(pinned memory) 등 운영 환경 최적화 정리
• Horace He, "Making Deep Learning Go Brrrr" — GPU 학습이 빠른 이유와 PyTorch 관점의 최적화 전략
• PyTorch Tensors 튜토리얼 — 텐서의 기본 연산을 한 페이지로 익힐 수 있는 공식 가이드
• PyTorch Autograd 튜토리얼 — requires_grad, .backward(), 계산 그래프 동작 방식을 직접 실험하며 익힘
• PyTorch nn.Module 레퍼런스 — nn.Module의 메서드, 등록 규칙, 후크(hook) 등의 공식 문서
• 모델 저장과 로딩 가이드 — state_dict 기반 저장/복원과 체크포인트 운영 패턴