PyTorch 기초 (3) - 모델 조립: nn.Module, 손실, 옵티마이저, 학습 루프

PyTorch 기초 시리즈(6편)의 3편입니다. 전체 목차는 0편에 있습니다. autograd로 기울기를 얻었다면, 이 편은 그 위에서 모델과 학습 과정을 조립합니다.

왜 nn.Module이 필요한가

앞 편에서는 파라미터 w 하나를 직접 만들고, 직접 갱신하고, 직접 zero_grad했습니다. 파라미터가 하나라 가능했습니다. 하지만 은닉층 몇 개짜리 신경망만 돼도 가중치와 편향이 수십, 수백 개가 됩니다. 이걸 일일이 리스트로 들고 다니며 갱신하는 것은 현실적이지 않습니다.

nn.Module은 이 파라미터들을 하나의 모델 객체 안에 캡슐화합니다. 모델을 정의하면 그 안의 모든 가중치가 자동으로 등록되고, model.parameters() 한 번으로 전부 꺼낼 수 있습니다.

모델 정의: nn.Module

모델은 nn.Module을 상속해 두 가지를 채우면 됩니다. __init__에서 층을 선언하고, forward에서 순전파를 기술합니다.

import torch.nn as nn

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(2, 8)   # 입력 2 → 은닉 8
        self.fc2 = nn.Linear(8, 1)   # 은닉 8 → 출력 1

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        return self.fc2(x)

model = Net()

nn.Linear(2, 8)은 딥러닝 기초 (1)에서 본 완전연결층 그 자체입니다 — 내부에 가중치 (8, 2)와 편향 (8,)을 requires_grad=True 상태로 들고 있습니다. 우리가 초기화하거나 추적할 필요 없이 알아서 관리됩니다.

층을 순서대로 쌓기만 하면 되는 단순한 경우에는 nn.Sequential이 더 간결합니다.

model = nn.Sequential(
    nn.Linear(2, 8),
    nn.ReLU(),
    nn.Linear(8, 1),
)

손실 함수와 옵티마이저

손실 함수는 nn이 표준적인 것들을 제공합니다. 회귀면 MSELoss, 분류면 CrossEntropyLoss를 주로 씁니다.

criterion = nn.MSELoss()           # 회귀
# criterion = nn.CrossEntropyLoss()  # 분류

옵티마이저는 (2)편에서 손으로 하던 w -= lr * w.grad를 대신합니다. torch.optim에서 골라, 모델의 파라미터를 넘겨 만듭니다. 딥러닝 기초 (5)에서 정리한 SGD·Adam이 여기 있습니다.

import torch.optim as optim

optimizer = optim.Adam(model.parameters(), lr=0.01)

model.parameters()가 열쇠입니다. 이 한 줄로 모델 안 모든 가중치가 옵티마이저에 연결되고, 이제 갱신은 옵티마이저가 통째로 처리합니다.

표준 학습 루프: 다섯 줄

여기까지 오면, 어떤 PyTorch 모델을 학습시키든 루프의 핵심은 똑같은 다섯 줄입니다. 이 형태를 손에 익히는 것이 이 편의 목표입니다.

for epoch in range(100):
    pred = model(X)              # 1. 순전파: 예측
    loss = criterion(pred, Y)    # 2. 손실 계산

    optimizer.zero_grad()        # 3. 이전 기울기 비우기
    loss.backward()              # 4. 역전파: 기울기 계산
    optimizer.step()             # 5. 갱신: 파라미터 한 걸음

각 줄이 앞에서 배운 것과 정확히 대응합니다.

단계	코드	하는 일
순전파	model(X)	forward를 실행해 예측을 냄
손실	criterion(pred, Y)	예측과 정답의 오차를 하나의 수로
초기화	optimizer.zero_grad()	누적된 기울기를 0으로 (2편의 그 함정)
역전파	loss.backward()	autograd가 모든 기울기를 계산
갱신	optimizer.step()	.grad를 보고 모든 파라미터를 갱신

zero_grad → backward → step의 순서가 특히 중요합니다. 순서가 틀리면 학습이 조용히 망가집니다. 이 세 줄을 하나의 관용구로 외워두면 됩니다.

정리

개념	한 줄 요약
nn.Module	파라미터를 캡슐화. __init__에 층, forward에 순전파
nn.Linear	완전연결층. 가중치·편향을 알아서 들고 추적
손실 함수	nn.MSELoss(회귀), nn.CrossEntropyLoss(분류)
옵티마이저	optim.Adam(model.parameters(), ...)가 갱신을 담당
학습 루프	순전파 → 손실 → zero_grad → backward → step

이제 모델을 정의하고 학습시킬 수 있습니다. 그런데 지금까지는 데이터 X, Y를 통째로 한 번에 넣었습니다. 실제로는 데이터가 메모리에 다 안 들어가거나, 미니배치로 나눠 섞어 먹여야 합니다. 다음 편은 그 데이터 공급을 담당하는 Dataset과 DataLoader입니다.

다음 글: PyTorch 기초 (4) - 데이터 파이프라인: Dataset과 DataLoader