본문 바로가기
TYLER SONGBlog
← Blog index
LLM Research9 min read

PyTorch Training Loop: zero_grad부터 optimizer.step까지 직접 쓰고 고치기

오늘은 신경망이 예측하고, 틀린 정도를 계산하고, 스스로 수정하는 전체 학습 루프를 직접 작성한다. 학습률과 기울기 초기화를 바꾸며 손실 변화를 관찰하고, 모양 불일치·NaN·메모리 부족 같은 오류를 원인부터 추적한다. 목표는 예제 없이도 학습 루프와 검증 루프를 백지에서 작성하는 것이다.

By 송민성

오늘의 목표

다음 순서를 보지 않고 작성할 수 있어야 한다.

zero_grad → forward → loss → backward → optimizer.step

MLP 학습 루프와 검증 루프를 직접 만들고, 일부러 넣은 세 가지 버그를 오류 메시지와 손실 기록을 이용해 고친다.

개념과 직관

학습 루프가 하는 일

신경망 학습을 과녁 맞히기에 비유해 보자.

  1. 이전 화살의 흔적을 지운다.
  2. 현재 자세로 화살을 쏜다.
  3. 과녁에서 얼마나 벗어났는지 잰다.
  4. 어느 방향으로 자세를 고쳐야 하는지 계산한다.
  5. 자세를 조금 수정한다.

PyTorch 코드와 연결하면 다음과 같다.

text
optimizer.zero_grad()  이전 기울기 제거
prediction = model(x)  순전파
loss = criterion(...)  손실 계산
loss.backward()        기울기 계산
optimizer.step()       파라미터 수정

순전파(forward pass, 입력을 신경망에 통과시켜 예측값을 만드는 과정)는 현재 파라미터로 답을 만든다.

손실(loss, 예측이 정답에서 얼마나 벗어났는지 나타내는 숫자)은 모델의 오답 정도다. 손실이 작을수록 현재 데이터에 더 잘 맞는다.

역전파(backpropagation, 손실을 줄이려면 각 파라미터를 어느 방향으로 바꿔야 하는지 계산하는 과정)를 실행하면 각 파라미터의 grad에 기울기가 저장된다.

기울기(gradient, 파라미터를 조금 바꿨을 때 손실이 어느 방향으로 얼마나 변하는지 나타내는 값)는 수정 방향을 알려주는 화살표다.

최소한의 수식

파라미터를 w, 학습률을 lr, 기울기를 grad라고 하면 기본 수정은 다음과 같다.

text
새로운 w = 기존 w - lr × grad

학습률(learning rate, 한 번에 파라미터를 얼마나 움직일지 정하는 값)이 너무 크면 최저점을 지나쳐 손실이 튄다. 너무 작으면 손실이 천천히 줄거나 거의 움직이지 않는다.

그림으로 보면 다음과 같다.

text
손실
 높음     ● 시작점
         / \
        /   \       lr이 너무 크면 골짜기를 계속 넘는다
       /     \  ● → ● → ●
      /       \
 낮음 --------●---------------- 파라미터
             최저점

왜 zero_grad()가 먼저인가

PyTorch는 backward()를 호출할 때 새 기울기를 기존 기울기에 더한다.

text
현재 grad = 이전 grad + 새로 계산한 grad

따라서 매 배치에서 zero_grad()를 호출하지 않으면 이전 배치의 기울기가 계속 쌓인다. 의도적으로 여러 배치의 기울기를 모으는 경우도 있지만, 오늘 만드는 일반 학습 루프에서는 매번 지워야 한다.

SGD와 Adam

옵티마이저(optimizer, 계산된 기울기를 이용해 파라미터를 수정하는 규칙)는 optimizer.step()에서 작동한다.

확률적 경사하강법(SGD, stochastic gradient descent, 현재 기울기에 학습률을 곱해 파라미터를 이동하는 방법)은 단순하고 움직임을 해석하기 쉽다.

text
w = w - lr × grad

Adam(Adaptive Moment Estimation, 최근 기울기의 방향과 크기를 함께 추적해 파라미터별 이동량을 조절하는 방법)은 파라미터마다 보폭을 조절한다. 복잡한 모델에서 초기 실험용으로 자주 쓰이지만, 어떤 문제에서나 자동으로 더 좋은 결과를 보장하지는 않는다.

python
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

또는:

python
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

검증 루프가 다른 이유

검증(validation, 학습에 사용하지 않은 데이터로 현재 모델의 성능을 확인하는 과정)에서는 파라미터를 수정하지 않는다.

python
model.eval() with torch.no_grad(): prediction = model(x) loss = criterion(prediction, y)

model.eval()은 드롭아웃처럼 학습과 평가 때 다르게 동작하는 층을 평가 상태로 바꾼다.

torch.no_grad()는 기울기 계산 기록을 만들지 않는다. 검증 속도가 빨라지고 메모리 사용량도 줄어든다. model.eval()과 torch.no_grad()는 역할이 다르므로 둘 다 사용한다.

코드로 직접 해보기

1단계: 라이브러리 없이 흐름 적기

코드를 작성하기 전에 종이에 다음 흐름을 적는다.

text
반복:
    이전 기울기를 지운다
    예측한다
    손실을 계산한다
    손실에서 각 파라미터까지 기울기를 계산한다
    파라미터를 수정한다

이 다섯 줄이 모든 기본 학습 루프의 뼈대다.

2단계: 실행 가능한 MLP 학습 코드

다층 퍼셉트론(MLP, multilayer perceptron, 선형 변환과 활성화 함수를 여러 층 쌓은 신경망)으로 두 종류의 점을 분류한다.

python
import torch from torch import nn from torch.utils.data import DataLoader, TensorDataset torch.manual_seed(42) # 두 입력값의 합이 0보다 크면 1, 아니면 0인 데이터 x = torch.randn(1000, 2) y = (x[:, 0] + x[:, 1] > 0).float() train_x, valid_x = x[:800], x[800:] train_y, valid_y = y[:800], y[800:] train_loader = DataLoader( TensorDataset(train_x, train_y), batch_size=32, shuffle=True, ) valid_loader = DataLoader( TensorDataset(valid_x, valid_y), batch_size=32, ) model = nn.Sequential( nn.Linear(2, 16), nn.ReLU(), nn.Linear(16, 1), ) criterion = nn.BCEWithLogitsLoss() optimizer = torch.optim.SGD(model.parameters(), lr=0.01) for epoch in range(20): model.train() train_loss_sum = 0.0 for batch_x, batch_y in train_loader: optimizer.zero_grad() logits = model(batch_x).squeeze(1) loss = criterion(logits, batch_y) loss.backward() optimizer.step() train_loss_sum += loss.item() * batch_x.size(0) model.eval() valid_loss_sum = 0.0 correct = 0 with torch.no_grad(): for batch_x, batch_y in valid_loader: logits = model(batch_x).squeeze(1) loss = criterion(logits, batch_y) valid_loss_sum += loss.item() * batch_x.size(0) predictions = (logits > 0).float() correct += (predictions == batch_y).sum().item() train_loss = train_loss_sum / len(train_loader.dataset) valid_loss = valid_loss_sum / len(valid_loader.dataset) valid_accuracy = correct / len(valid_loader.dataset) print( f"에폭 {epoch + 1:02d} | " f"학습 손실 {train_loss:.4f} | " f"검증 손실 {valid_loss:.4f} | " f"검증 정확도 {valid_accuracy:.3f}" )

BCEWithLogitsLoss는 이진 분류용 손실 함수다. 모델의 마지막 출력에 시그모이드를 직접 붙이지 않고 원래 점수인 로짓(logit, 확률로 바꾸기 전의 실수 점수)을 그대로 전달한다.

모델 출력 모양은 [배치 크기, 1]이다. 정답 모양은 [배치 크기]이므로 .squeeze(1)로 출력의 마지막 크기 1을 제거한다.

3단계: 일부러 망가뜨린 코드 고치기

버그 1: zero_grad() 누락

망가진 부분:

python
for batch_x, batch_y in train_loader: logits = model(batch_x).squeeze(1) loss = criterion(logits, batch_y) loss.backward() optimizer.step()

문제는 매 배치의 기울기가 이전 기울기에 누적된다는 점이다. 손실이 불안정해지거나 갑자기 커질 수 있다.

수정:

python
for batch_x, batch_y in train_loader: optimizer.zero_grad() logits = model(batch_x).squeeze(1) loss = criterion(logits, batch_y) loss.backward() optimizer.step()

버그 2: 학습률 100

망가진 부분:

python
optimizer = torch.optim.SGD(model.parameters(), lr=100)

한 번의 수정 폭이 너무 커서 손실이 급증하거나 nan이 나타날 수 있다.

수정:

python
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

학습률 문제를 의심할 때는 현재 값보다 10배 작은 값을 먼저 시험한다.

버그 3: 모양 불일치

망가진 부분:

python
logits = model(batch_x) # [32, 1] loss = criterion(logits, batch_y) # [32]

예측과 정답의 모양이 다르다.

수정:

python
logits = model(batch_x).squeeze(1) # [32] loss = criterion(logits, batch_y) # [32]

오류를 고치기 전에 실제 모양부터 출력한다.

python
print("예측 모양:", logits.shape) print("정답 모양:", batch_y.shape)

오늘의 실험

실험 1: 학습률 비교

다음 세 값을 각각 사용해 새 모델을 처음부터 20에폭 학습한다.

python
learning_rates = [0.1, 0.01, 0.001]

학습률마다 같은 초기값을 사용해야 비교가 쉬워진다.

python
for lr in learning_rates: torch.manual_seed(42) model = nn.Sequential( nn.Linear(2, 16), nn.ReLU(), nn.Linear(16, 1), ) optimizer = torch.optim.SGD(model.parameters(), lr=lr)

각 에폭의 학습 손실과 검증 손실을 표에 기록한다.

학습률1에폭 학습 손실10에폭 학습 손실20에폭 학습 손실20에폭 검증 손실관찰
0.1
0.01
0.001

관찰할 질문:

  1. 초반 손실이 가장 빠르게 줄어드는 값은 무엇인가?
  2. 손실이 진동하거나 갑자기 커지는 값이 있는가?
  3. 20에폭 뒤에도 너무 천천히 줄어드는 값이 있는가?
  4. 학습 손실과 검증 손실이 함께 줄어드는가?

학습률의 결과는 모델, 데이터, 옵티마이저에 따라 달라진다. 숫자를 외우지 말고 손실 곡선의 움직임을 읽는다.

실험 2: zero_grad() 제거

같은 모델과 학습률로 두 번 실행한다.

  • 실행 A: optimizer.zero_grad() 사용
  • 실행 B: optimizer.zero_grad() 제거

다른 조건은 모두 같게 둔다.

에폭정상 학습 손실zero_grad 제거 손실
1
5
10
20

zero_grad()를 제거했을 때 손실이 반드시 같은 방식으로 망가지는 것은 아니다. 데이터 순서와 초기값에 따라 잠시 줄어들 수도 있다. 핵심 관찰은 각 배치의 기울기가 독립적으로 사용되지 않고 계속 누적된다는 점이다.

기울기 크기도 확인할 수 있다.

python
first_layer = model[0] print(first_layer.weight.grad.norm().item())

grad.norm()은 기울기 전체 크기를 하나의 숫자로 요약한다. zero_grad() 유무에 따라 이 값이 어떻게 변하는지 기록한다.

왜 이걸 배우나 (LLM / Post-training 연결)

LLM 사후학습(post-training, 사전학습된 언어 모델을 특정 행동과 목적에 맞게 추가 학습하는 과정)도 기본 구조는 같은 학습 루프다. 입력 문장으로 순전파하고, 목표 출력과 비교해 손실을 계산하고, 역전파한 뒤 옵티마이저로 파라미터를 수정한다.

실무에서는 거대한 모델, 여러 GPU, 기울기 누적, 혼합 정밀도 같은 요소가 추가된다. 손실 급증·NaN·메모리 부족이 발생했을 때 기본 학습 루프를 이해해야 데이터 문제인지, 학습률 문제인지, 기울기 문제인지 구분할 수 있다.

초보자가 막히는 지점

손실이 줄지 않는다

먼저 한두 개 배치만 반복 학습해 본다. 작은 데이터조차 외우지 못하면 전체 데이터보다 코드나 설정 문제일 가능성이 크다.

확인 순서:

  1. loss.backward()와 optimizer.step()이 실행되는가?
  2. 옵티마이저에 model.parameters()를 전달했는가?
  3. 학습률이 너무 작거나 큰가?
  4. 입력과 정답이 올바르게 짝지어졌는가?
  5. 손실 함수가 현재 문제와 맞는가?

기울기가 None이다

python
print(model[0].weight.grad)

None이라면 다음 원인을 확인한다.

  • loss.backward() 전에 출력했는가?
  • 순전파를 torch.no_grad() 안에서 실행했는가?
  • 중간 결과에 .detach()를 사용했는가?
  • 해당 파라미터가 손실 계산에 실제로 사용됐는가?
  • 파라미터의 requires_grad가 False인가?

모양 불일치 오류가 난다

계산 직전에 모양을 출력한다.

python
print(batch_x.shape) print(logits.shape) print(batch_y.shape)

reshape나 squeeze를 무작정 추가하지 않는다. 각 축이 배치, 특징, 클래스 중 무엇을 뜻하는지 먼저 적는다.

손실이 nan이 된다

NaN(Not a Number, 계산할 수 없는 숫자 상태)은 보통 값이 지나치게 커졌거나 잘못된 연산이 들어갔다는 신호다.

python
print("입력 정상:", torch.isfinite(batch_x).all().item()) print("출력 정상:", torch.isfinite(logits).all().item()) print("손실 정상:", torch.isfinite(loss).item())

학습률을 10배 낮추고, 어느 단계에서 처음 유한하지 않은 값이 생기는지 찾는다.

메모리 부족 오류가 난다

OOM(Out Of Memory, GPU나 시스템 메모리가 부족한 상태)이 발생하면 배치 크기를 먼저 줄인다.

python
train_loader = DataLoader( TensorDataset(train_x, train_y), batch_size=16, shuffle=True, )

검증에서는 torch.no_grad()를 사용한다. 에폭마다 손실 텐서 자체를 리스트에 저장하지 말고 loss.item()으로 숫자만 저장한다.

python
# 피하기 losses.append(loss) # 사용하기 losses.append(loss.item())

알고리즘 (병행 1시간)

쉬움 1문제: 두 수의 합

정수 배열과 목표값이 주어졌을 때 합이 목표값이 되는 두 원소의 인덱스를 찾는다.

권장 시간:

  • 문제 이해와 예시 작성: 5분
  • 직접 구현: 15분
  • 더 빠른 접근 찾기: 10분

먼저 모든 두 원소 조합을 확인하는 풀이를 작성한다. 이후 딕셔너리에 이전 숫자와 인덱스를 저장해 한 번의 순회로 푼다.

기록할 것:

  • 처음 생각한 접근
  • 최종 접근
  • 시간복잡도와 공간복잡도
  • 틀렸다면 틀린 입력과 이유

중간 1문제: 나선형 행렬

행렬의 바깥쪽부터 시계 방향으로 돌며 모든 원소를 순서대로 반환한다.

권장 시간:

  • 이동 규칙 그림 그리기: 10분
  • 경계 변수로 구현: 20분
  • 한 줄·한 열 입력 점검: 15분

위, 아래, 왼쪽, 오른쪽 경계를 두고 한 방향을 순회할 때마다 해당 경계를 안쪽으로 옮긴다. 아래쪽이나 왼쪽 방향을 읽기 전에는 아직 읽지 않은 행과 열이 남았는지 확인한다.

기록할 것:

  • 네 방향을 어떤 순서로 처리했는지
  • 각 경계가 언제 변하는지
  • 시간복잡도와 공간복잡도
  • 중복 방문이 생겼다면 어떤 경계 조건이 틀렸는지