학습 루프 직접 작성하기

학습 루프의 큰 그림

이제 지금까지 배운 조각을 하나로 연결할 차례입니다. 모델은 입력을 받아 예측값을 만들고, 손실 함수는 예측이 정답과 얼마나 다른지 계산합니다. 옵티마이저는 그 손실을 줄이는 방향으로 모델의 파라미터를 조금씩 바꿉니다.

이 과정을 한 번만 하면 학습이 거의 일어나지 않습니다. 그래서 전체 데이터를 여러 번 반복해서 봅니다. 전체 데이터셋을 한 번 모두 사용하는 단위를 epoch라고 합니다. 그리고 DataLoader가 나누어 주는 작은 데이터 묶음을 batch라고 합니다.

기본 구조는 다음과 같습니다.

epochs = 5

for epoch in range(epochs):
    print(f'epoch {epoch + 1}')

    for batch, (X, y) in enumerate(train_dataloader):
        # 한 배치 학습
        pass

바깥쪽 for epoch는 같은 데이터셋을 몇 번 반복해서 볼지 정합니다. 안쪽 for batch는 한 epoch 안에서 배치를 하나씩 꺼냅니다. 실제 학습은 대부분 이 안쪽 반복문에서 일어납니다.

한 배치에서 일어나는 일

학습을 시작할 때는 먼저 모델을 학습 모드로 바꿉니다.

model.train()

model.train()은 모델에게 "지금은 학습 중"이라고 알려 주는 역할을 합니다. 단순한 신경망에서는 차이가 잘 보이지 않을 수 있지만, 드롭아웃이나 배치 정규화처럼 학습과 평가 때 다르게 동작하는 층이 있으면 반드시 필요합니다.

한 배치의 학습 코드는 다음 순서로 작성합니다.

X = X.to(device)
y = y.to(device)

pred = model(X)
loss = loss_fn(pred, y)

optimizer.zero_grad()
loss.backward()
optimizer.step()

먼저 입력 X와 정답 y를 모델이 있는 장치로 보냅니다. 모델이 GPU에 있는데 데이터가 CPU에 있으면 계산할 수 없기 때문입니다.

그다음 pred = model(X)로 예측값을 만듭니다. 이것을 순전파, 또는 forward라고 부릅니다. 이어서 loss_fn(pred, y)로 예측값과 정답의 차이를 손실값으로 계산합니다.

세 줄이 학습을 만든다

PyTorch 학습 루프에서 가장 중요한 세 줄은 다음입니다.

optimizer.zero_grad()
loss.backward()
optimizer.step()

optimizer.zero_grad()는 이전 배치에서 계산된 기울기를 지웁니다. PyTorch는 기본적으로 기울기를 누적합니다. 따라서 매 배치마다 지우지 않으면 이전 계산이 섞여 원하지 않는 방향으로 학습될 수 있습니다.

loss.backward()는 손실을 줄이기 위해 각 파라미터를 어느 방향으로 바꾸어야 하는지 계산합니다. 이 값은 각 파라미터의 .grad에 저장됩니다.

optimizer.step()은 .grad에 저장된 값을 사용해 실제로 파라미터를 갱신합니다. 즉, 이 줄에서 모델의 가중치가 바뀝니다.

순서는 반드시 지키는 것이 좋습니다. 기울기를 지우고, 새 기울기를 계산하고, 그 기울기로 파라미터를 바꿉니다.

손실값 출력하기

학습이 진행되는지 확인하려면 중간중간 손실값을 출력하면 됩니다. 텐서 형태의 손실값을 파이썬 숫자로 꺼낼 때는 loss.item()을 사용합니다.

if batch % 100 == 0:
    loss_value = loss.item()
    print(f'batch {batch}, loss: {loss_value:.4f}')

손실값이 항상 매끄럽게 내려가지는 않습니다. 배치마다 데이터가 다르기 때문입니다. 하지만 여러 epoch를 지나면서 전체적으로 줄어드는 흐름이 보이면 학습이 되고 있다고 볼 수 있습니다.

train_loop 함수로 분리하기

학습 코드를 매번 직접 쓰면 중복이 많아집니다. 그래서 보통 학습 절차를 함수로 분리합니다.

def train_loop(dataloader, model, loss_fn, optimizer, device):
    model.train()

    for batch, (X, y) in enumerate(dataloader):
        X = X.to(device)
        y = y.to(device)

        pred = model(X)
        loss = loss_fn(pred, y)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if batch % 100 == 0:
            print(f'batch {batch}, loss: {loss.item():.4f}')

이제 전체 학습은 더 단순해집니다.

epochs = 5

for epoch in range(epochs):
    print(f'epoch {epoch + 1}')
    train_loop(train_dataloader, model, loss_fn, optimizer, device)

이 구조가 PyTorch 학습 코드의 중심입니다. PyTorch는 학습 과정을 숨기지 않습니다. 대신 사용자가 forward, loss 계산, backward, optimizer step을 직접 연결하게 합니다. 처음에는 코드가 길어 보이지만, 이 흐름을 이해하면 어떤 모델을 학습하더라도 같은 뼈대를 사용할 수 있습니다.