idlemachines의 The annotated PyTorch training loop는 PyTorch 학습 루프를 한 줄씩 해부하는 글입니다. 겉으로는 단순한 for 문처럼 보이지만, 실제로는 모델 상태, gradient buffer, autograd graph, optimizer state, scheduler state가 맞물려 돌아갑니다. 그래서 줄 하나가 조금만 엉뚱한 곳에 있어도 에러 없이 학습이 망가질 수 있습니다.

원문: The annotated PyTorch training loop

1) 학습 루프는 작은 상태 머신이다

PyTorch 학습 코드는 보통 이렇게 생겼습니다.

for epoch in range(num_epochs):
    model.train()

    for x_batch, y_batch in loader:
        optimiser.zero_grad()
        logits = model(x_batch)
        loss = criterion(logits, y_batch)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        optimiser.step()

    scheduler.step()

    model.eval()
    with torch.no_grad():
        val_logits = model(x_val)
        val_loss = criterion(val_logits, y_val)

중요한 점은 이 코드가 단순한 실행 순서가 아니라는 것입니다. 각 줄은 다음 상태 중 하나를 바꿉니다.

줄바뀌는 상태
model.train()Dropout, BatchNorm 같은 layer의 동작 모드
zero_grad()parameter마다 쌓여 있는 .grad buffer
model(x)autograd graph와 activation 저장
loss.backward()각 parameter의 .grad 값
clip_grad_norm_()이미 계산된 gradient의 크기
optimiser.step()실제 model weight와 optimizer 내부 상태
scheduler.step()optimizer의 learning rate
model.eval()검증/추론용 layer 동작
torch.no_grad()graph 생성 여부

그래서 학습 루프를 이해한다는 것은 “이 줄이 무슨 계산을 하는가"보다 “이 줄이 어떤 상태를 바꾸는가"를 아는 일에 가깝습니다.

2) 조용히 망가지는 실수가 많다

원문이 특히 강조하는 부분은, 많은 실수가 즉시 터지지 않는다는 점입니다. 코드는 돌아가지만 결과가 이상해집니다.

실수겉으로 보이는 현상
optimiser.zero_grad()를 loss.backward() 뒤에 둠이전 batch의 gradient가 더해져 현재 batch만의 update가 아니게 됨
clip_grad_norm_()을 backward() 전에 호출.grad가 비어 있어서 사실상 아무 일도 하지 않음
clip_grad_norm_()을 optimiser.step() 뒤에 호출이미 적용된 gradient를 뒤늦게 자르므로 효과가 없음
scheduler.step()을 batch loop 안에 둠epoch마다 한 번 줄어야 할 learning rate가 batch 수만큼 빨리 줄어듦
검증 후 model.train()으로 돌아가지 않음Dropout이 꺼지고 BatchNorm이 고정된 상태로 학습됨
검증에서 torch.no_grad()를 빼먹음validation마다 graph가 생겨 메모리를 잡아먹음
logging에 loss tensor를 그대로 저장graph 참조가 남아 메모리가 계속 붙잡힘

이런 문제는 NameError처럼 친절하게 알려주지 않습니다. loss가 이상하게 흔들리거나, 학습이 안 되거나, GPU 메모리가 점점 차는 식으로 나타납니다. 그래서 루프의 순서는 암기 대상이 아니라 디버깅 체크리스트입니다.

3) zero_grad()는 왜 매 batch마다 필요할까

PyTorch의 .grad는 덮어쓰기 방식이 아니라 누적 방식입니다. loss.backward()를 호출하면 새 gradient가 기존 .grad에 더해집니다.

이 설계 덕분에 gradient accumulation이 가능합니다. GPU 메모리가 부족해서 큰 batch를 한 번에 못 올릴 때, 작은 batch 여러 개의 gradient를 모은 뒤 한 번만 optimiser.step()을 호출할 수 있습니다.

하지만 일반 학습에서는 batch마다 독립적인 update를 원합니다. 그래서 batch 시작 시점에 optimiser.zero_grad()를 호출해 이전 batch의 gradient를 지웁니다.

실무에서는 zero_grad(set_to_none=True)도 자주 씁니다. 0으로 채우는 대신 .grad를 None으로 두기 때문에 약간 더 빠르고 메모리 사용도 줄일 수 있습니다. PyTorch 2.0 이후 기본 동작도 이 방향입니다.

4) forward는 graph를 만들고, loss는 출발점을 만든다

logits = model(x_batch)는 단순히 예측값을 얻는 줄이 아닙니다. requires_grad=True인 tensor가 연산에 참여하면 PyTorch는 그 연산들을 동적으로 기록합니다. 이 기록이 나중에 backward()가 따라갈 computation graph입니다.

이 graph에는 backward 때 필요한 activation도 함께 살아 있습니다. 그래서 학습 중 forward는 메모리를 많이 씁니다. 반대로 검증에서는 gradient가 필요 없으므로 torch.no_grad() 또는 torch.inference_mode()를 사용해 graph 생성을 막습니다.

분류 문제에서 흔히 쓰는 nn.CrossEntropyLoss는 raw logits를 받습니다. 내부적으로 log-softmax와 negative log-likelihood를 결합해 계산합니다. 따라서 모델 끝에 softmax를 붙여 확률로 만든 뒤 넣는 방식은 보통 피해야 합니다. 수치 안정성도 떨어지고, 의도한 API 사용법도 아닙니다.

또 하나의 작은 함정은 logging입니다. 화면에 찍거나 리스트에 저장할 값은 loss.item()으로 Python 숫자만 뽑아야 합니다. tensor 자체를 오래 붙잡으면 graph 참조도 함께 남아 메모리 문제가 생길 수 있습니다.

5) backward()와 step()은 다른 일이다

초보자가 가장 헷갈리기 쉬운 지점입니다.

loss.backward()는 weight를 바꾸지 않습니다. scalar loss에서 시작해 graph를 거꾸로 타고 내려가며 각 parameter의 .grad를 채웁니다. 이것은 “어느 방향으로 얼마나 바꾸면 loss가 줄어드는가"를 계산하는 단계입니다.

실제 weight 변경은 optimiser.step()에서만 일어납니다. optimizer는 .grad를 읽고, SGD나 Adam 같은 규칙에 따라 parameter를 갱신합니다. Adam이라면 gradient 평균과 제곱 평균 같은 내부 moment state도 함께 갱신합니다.

이 둘 사이에 gradient clipping을 넣는 이유도 여기 있습니다.

  1. backward() 전에는 .grad가 아직 없습니다.
  2. backward() 후에는 .grad가 계산되어 있습니다.
  3. step() 전에는 아직 weight가 바뀌지 않았습니다.

따라서 clipping의 자연스러운 위치는 loss.backward() 다음, optimiser.step() 이전입니다.

6) scheduler는 종류마다 호출 위치가 다르다

기본 예시에서는 scheduler.step()을 epoch 끝에서 한 번 호출합니다. 예를 들어 cosine annealing을 epoch 단위로 설계했다면 batch loop 밖에 있어야 합니다.

하지만 모든 scheduler가 같은 규칙을 따르지는 않습니다. Hugging Face의 warmup + cosine scheduler처럼 step 단위로 설계된 scheduler는 batch마다 호출합니다. ReduceLROnPlateau는 예외적으로 validation loss 같은 metric을 받아야 하므로 scheduler.step(val_loss)처럼 호출합니다.

즉 scheduler의 핵심 질문은 “이 learning rate 계획이 epoch 단위인가, update step 단위인가, metric 기반인가"입니다. 이 질문 없이 위치만 복사하면 learning rate가 의도보다 훨씬 빨리 줄어들 수 있습니다.

7) 검증에는 eval()과 no_grad()가 둘 다 필요하다

model.eval()과 torch.no_grad()는 이름이 비슷한 역할처럼 보이지만 완전히 다릅니다.

model.eval()은 module의 동작 모드를 바꿉니다. Dropout은 mask를 샘플링하지 않고 그대로 통과시키며, BatchNorm은 현재 batch 통계가 아니라 학습 중 쌓아둔 running statistics를 사용합니다.

torch.no_grad()는 autograd graph 생성을 끕니다. 출력 tensor에 grad_fn이 붙지 않고, backward를 위한 activation도 저장하지 않습니다. 그래서 검증이 더 빠르고 메모리도 덜 씁니다.

두 장치는 독립적입니다.

조합사용 사례
eval() + graph enabledsaliency map, adversarial example처럼 입력 gradient가 필요한 분석
train() + no_grad()학습 모드 forward 동작만 빠르게 확인
eval() + no_grad()일반적인 validation 또는 inference

최근 PyTorch에서는 검증/추론에 torch.inference_mode()를 쓰기도 합니다. no_grad()보다 더 엄격하고 약간 더 빠를 수 있지만, 그 안에서 만들어진 tensor를 나중에 gradient 계산에 섞기 어렵다는 점을 알고 써야 합니다.

8) checkpoint는 model만 저장하면 부족하다

학습을 중단했다가 이어서 돌릴 가능성이 있다면 model weight만 저장해서는 부족합니다. optimizer state와 scheduler state도 함께 저장해야 합니다.

특히 Adam은 parameter마다 moment estimate를 들고 있습니다. 이 상태 없이 다시 시작하면 weight는 이어받았더라도 optimizer는 처음부터 워밍업하는 셈이 되어 loss가 튈 수 있습니다.

일반적으로는 다음 정보를 checkpoint에 넣습니다.

  • epoch
  • model.state_dict()
  • optimiser.state_dict()
  • scheduler.state_dict()
  • validation loss 또는 best score

그리고 최종 epoch의 weight보다 validation 성능이 가장 좋았던 시점의 weight를 따로 저장하는 패턴이 많습니다. 실제 배포나 평가에는 마지막 model보다 가장 일반화가 잘 된 model이 더 유용하기 때문입니다.

9) GPU 효율은 루프 바깥 설정과 루프 안 이동이 함께 만든다

GPU 학습에서는 모델과 batch가 같은 device에 있어야 합니다. 그리고 optimizer는 model을 device로 옮긴 뒤에 만들어야 합니다. dtype 변환이나 device 이동이 parameter object를 새로 만들 수 있기 때문에, optimizer가 예전 parameter를 붙잡는 상황을 피해야 합니다.

batch 이동에서는 non_blocking=True를 자주 봅니다.

x_batch = x_batch.to(device, non_blocking=True)
y_batch = y_batch.to(device, non_blocking=True)

이 옵션은 host-to-device 전송을 비동기로 시작하게 해줍니다. 단, DataLoader 쪽에서 pin_memory=True를 사용해야 효과가 있습니다.

mixed precision도 중요한 무료 속도 향상 수단입니다. float16이나 bfloat16 연산은 최신 GPU tensor core에서 훨씬 빠르게 처리됩니다. 다만 float16은 표현 범위가 좁아 작은 gradient가 0으로 사라질 수 있으므로 GradScaler를 함께 사용합니다. bfloat16은 float32와 exponent 범위가 같아 보통 scaler가 필요 없습니다.

DataLoader도 병목이 될 수 있습니다. num_workers, pin_memory, persistent_workers, prefetch_factor를 잘 조정하면 GPU가 data를 기다리며 쉬는 시간을 줄일 수 있습니다.

마지막으로 PyTorch 2.0 이상의 torch.compile은 forward를 추적해 연산을 fusion하고 더 효율적인 kernel을 만들 수 있습니다. 첫 실행은 compile 때문에 느릴 수 있지만, 긴 학습에서는 그 비용이 상쇄됩니다.

결론 및 시사점

PyTorch 학습 루프의 어려움은 문법보다 상태 관리에 있습니다. zero_grad()는 gradient buffer를 다루고, forward는 graph를 만들고, backward()는 gradient를 채우고, step()은 weight를 바꾸고, eval()과 no_grad()는 검증 시점의 동작과 메모리 사용을 제어합니다.

이 구분이 선명해지면 학습이 안 될 때 확인할 질문도 분명해집니다. gradient가 이전 batch와 섞였는가? clipping이 실제로 적용되었는가? scheduler가 너무 자주 호출되었는가? 검증에서 graph를 만들고 있지는 않은가? optimizer가 올바른 parameter를 보고 있는가?

한 줄 결론: PyTorch 학습 루프는 외워야 할 주문이 아니라, 모델 상태와 gradient 흐름을 올바른 순서로 통제하는 작은 실행 시스템이다.