17.5 학습 루프 해부: forward → loss → backward → update

실습 · 연습문제 2개

gradient accumulation이 있는 학습 루프 accumulation-loop

메모리 때문에 큰 batch를 한 번에 올릴 수 없으면 작은 micro-batch 여러 개의 gradient를 모아 한 번 갱신한다(10.4장). 올바르게 구현하면 micro-batch 크기 8 × 누적 4회는 batch 크기 32와 같은 학습이어야 한다. 선형회귀 y^=x⊤w+b\hat{y} = \mathbf{x}^\top\mathbf{w} + b, 손실 L=1n∑i(y^i−yi)2L = \frac{1}{n}\sum_i(\hat{y}_i - y_i)^2의 학습 루프를 NumPy로 쓴다.

train_linear(X, y, micro_batch, accum, lr, epochs, seed)를 완성한다. 반환은 (w, b, n_steps)이고, n_steps는 파라미터를 갱신한 횟수(optimizer step 수)다.

규칙은 다음과 같다.

  1. 파라미터는 w = 0(shape (D,)), b = 0에서 시작한다.
  2. epoch e의 표본 순서는 np.random.default_rng([seed, e]).permutation(N)이다.
  3. 그 순서를 앞에서부터 micro_batch개씩 자른다(마지막 micro-batch는 작을 수 있다).
  4. micro-batch accum개마다, 그리고 epoch이 끝날 때 남은 micro-batch가 있으면, 모은 gradient로 SGD 한 step을 한다: θ←θ−lr⋅g\theta \leftarrow \theta - \text{lr}\cdot g.
  5. gg는 그 step에 쓰인 모든 표본에 대한 평균 손실의 gradient여야 한다. micro-batch마다 크기가 다를 수 있으므로, micro-batch 평균 gradient를 단순 평균하지 말고 표본 수로 가중한다(또는 표본별 gradient의 합을 모아 마지막에 총 표본 수로 나눈다).
  6. gradient를 모으는 동안에는 파라미터를 바꾸지 않는다. step이 끝나면 모은 gradient를 비운다(zero_grad).

테스트는 (micro 8, accum 4)와 (micro 32, accum 1)의 결과가 같은지, step 수가 맞는지, 표본 수가 나누어떨어지지 않을 때도 같은지, 학습이 실제로 문제를 푸는지 확인한다.

`.grad` 버퍼 시뮬레이터로 학습 루프의 순서 버그 찾기 grad-buffer-sim

PyTorch 학습 루프의 세 호출 optimizer.zero_grad(), loss.backward(), optimizer.step()은 모두 파라미터마다 붙어 있는 .grad 버퍼를 읽거나 쓴다. 순서를 바꾸거나 하나를 빠뜨려도 오류는 나지 않는다. 버퍼의 규칙만 알면 결과를 예측할 수 있다. 파라미터가 스칼라 하나인 경우의 시뮬레이터를 만든다.

run_ops(ops, w0, lr, grad_fn)을 완성한다.

  • 상태: 파라미터 w(처음 w0)와 버퍼 grad(처음 None).
  • ops의 문자열을 차례로 실행한다.
    • "zero_grad": grad = None (집필 시점 PyTorch 기본값 set_to_none=True와 같다)
    • "backward": g = grad_fn(w)를 계산해 버퍼에 더한다. 버퍼가 None이면 g가 된다.
    • "step": 버퍼가 None이면 아무것도 하지 않는다. 아니면 w = w - lr * grad (SGD). 버퍼는 그대로 둔다.
    • 그 밖의 문자열이면 ValueError.
  • 반환: (w, grad) — 모든 연산 뒤의 파라미터와 버퍼.

backward는 현재 w에서 gradient를 계산한다는 점에 주의한다. 테스트는 다음 경우를 확인한다.

한 반복의 순서 결과
zero_grad, backward, step 올바른 SGD
backward, step (zero_grad 누락) 과거 gradient가 계속 쌓인다
backward, zero_grad, step 버퍼가 비어 파라미터가 전혀 바뀌지 않는다
zero_grad, backward, backward, step gradient accumulation: 두 gradient의 합으로 한 번 갱신