Trainer 설계: 직접 작성과 프레임워크

trainer.py의 책임

trainer가 무엇을 하는지보다 무엇을 하지 않는지를 먼저 정하는 것이 설계의 핵심이다. 6.1장에서 책임의 경계를 "데이터가 넘어가는 지점"으로 정했던 것처럼, trainer와 나머지 코드의 경계도 주고받는 것으로 정의한다.

왼쪽은 실험이 바뀌어도 바뀌지 않아야 하는 부분이다. 여기를 한곳에 모으면 "두 실험이 루프의 세부에서 달랐다"는 가능성이 사라진다. 오른쪽은 연구 질문 그 자체이므로 trainer 밖에 둔다. 경계를 넘는 것은 함수 두 개다.

• step 함수 step_fn(model, batch) -> loss: 과제별 forward와 loss.
• 평가 함수 eval_fn(model) -> dict: 평가 모드 전환, 합/개수 누적, metric 계산(17.6장).

이 경계를 NumPy로 먼저 그려 본다. 아래 Trainer는 반복문의 순서와 이벤트 발송만 알고, 선형 회귀든 무엇이든 손실과 업데이트는 step_fn이 정한다.

import numpy as np

class Trainer:
    def __init__(self, step_fn, eval_fn, callbacks=(), max_epochs=100):
        self.step_fn, self.eval_fn = step_fn, eval_fn          # 과제별 코드는 밖에서 받는다
        self.callbacks, self.max_epochs = list(callbacks), max_epochs
        self.should_stop, self.epoch = False, 0

    def _emit(self, event, **info):
        for cb in self.callbacks:
            getattr(cb, event, lambda trainer, **kw: None)(self, **info)

    def fit(self, params, batches):
        for self.epoch in range(1, self.max_epochs + 1):
            for batch in batches:
                params, loss = self.step_fn(params, batch)     # forward·backward·update는 과제가 정한다
            self._emit("on_epoch_end", train_loss=loss)
            self._emit("on_eval_end", metric=self.eval_fn(params))
            if self.should_stop:
                break
        return params

부가 기능은 callback으로 붙인다. callback은 trainer가 정해진 시점(이벤트)에 호출하는 객체다. trainer는 callback이 무엇을 하는지 모르고, callback은 trainer의 공개된 상태(epoch, should_stop)만 읽고 쓴다. 아래는 early stopping과 기록을 callback으로 만들고, 같은 trainer로 손실이 다른 두 과제(MSE, Huber)를 학습한다. 잡음이 두꺼운 꼬리(자유도 2의 t 분포)를 가져서 이상치에 덜 민감한 Huber 손실이 의미가 있는 상황이다.

class EarlyStopping:
    def __init__(self, patience, min_delta=1e-3):
        self.patience, self.min_delta, self.best, self.wait = patience, min_delta, np.inf, 0
    def on_eval_end(self, trainer, metric):
        improved = metric < self.best - self.min_delta
        self.best, self.wait = (metric, 0) if improved else (self.best, self.wait + 1)
        trainer.should_stop = self.wait >= self.patience

class History:
    def __init__(self):
        self.rows = []
    def on_eval_end(self, trainer, metric):
        self.rows.append((trainer.epoch, round(float(metric), 4)))

rng = np.random.default_rng(0)
X = rng.normal(size=(200, 3))                                  # (N, D)
y = X @ np.array([1.0, -2.0, 0.5]) + rng.standard_t(df=2, size=200) * 0.3   # 두꺼운 꼬리 잡음
batches = [(X[i:i + 20], y[i:i + 20]) for i in range(0, 160, 20)]
Xv, yv = X[160:], y[160:]

def mse_step(w, b, lr=0.05):
    r = b[0] @ w - b[1]
    return w - lr * 2 * b[0].T @ r / len(r), np.mean(r ** 2)
def huber_step(w, b, lr=0.05, delta=1.0):
    r = b[0] @ w - b[1]
    return w - lr * b[0].T @ np.clip(r, -delta, delta) / len(r), np.mean(np.minimum(r ** 2 / 2, delta * (np.abs(r) - delta / 2)))
val_mae = lambda w: np.mean(np.abs(Xv @ w - yv))

for name, step in [("MSE", mse_step), ("Huber", huber_step)]:
    hist = History()
    w = Trainer(step, val_mae, [EarlyStopping(patience=3), hist], max_epochs=50).fit(np.zeros(3), batches)
    print(f"{name:5s} 멈춘 epoch {hist.rows[-1][0]:2d}  검증 MAE {hist.rows[-1][1]:.4f}  w = {np.round(w, 3)}")

MSE   멈춘 epoch 11  검증 MAE 0.3788  w = [ 1.08  -1.945  0.478]
Huber 멈춘 epoch 22  검증 MAE 0.3747  w = [ 0.972 -1.966  0.498]

Trainer 클래스는 한 줄도 바뀌지 않았다. 두 과제의 차이는 step 함수에만 있고, 멈춘 시점은 각 과제의 검증 곡선이 정했다. 이것이 재사용 가능한 trainer의 기준이다. 새 과제를 추가할 때 trainer 파일을 열 필요가 없어야 한다.

PyTorch에서도 구조는 같다. 다만 이제 backward와 optimizer.step()의 순서가 trainer의 책임으로 들어온다. 이 순서는 17.5장에서 본 대로 틀리기 쉽고, 모든 과제에서 같기 때문이다. trainer로 옮긴 코드는 옮기기 전과 같은 결과를 내야 한다. 리팩터링이 결과를 바꾸지 않았다는 확인은 파라미터 비교로 한다.

import torch
from torch import nn

def fit(model, opt, loader, step_fn, max_epochs):
    for epoch in range(max_epochs):
        model.train()
        for batch in loader:
            loss = step_fn(model, batch)                       # 과제별: forward와 loss
            opt.zero_grad()
            loss.backward()
            opt.step()
    return model

def regression_step(model, batch):
    x, y = batch
    return nn.functional.mse_loss(model(x), y)

def plain_loop(model, opt, loader, max_epochs):                # 17.5장의 루프를 그대로 쓴 것
    for epoch in range(max_epochs):
        model.train()
        for x, y in loader:
            loss = nn.functional.mse_loss(model(x), y)
            opt.zero_grad()
            loss.backward()
            opt.step()
    return model

def setup(seed=0):
    torch.manual_seed(seed)
    x = torch.randn(64, 3)                                     # (N, 3)
    y = torch.sin(x[:, :1]) + x[:, 1:2] * x[:, 2:]              # (N, 1)
    data = list(zip(x.split(16), y.split(16)))                 # batch 4개
    model = nn.Sequential(nn.Linear(3, 16), nn.Tanh(), nn.Linear(16, 1))
    return model, torch.optim.Adam(model.parameters(), lr=1e-2), data

a = fit(*setup(), regression_step, max_epochs=200)
b = plain_loop(*setup(), max_epochs=200)
print("trainer와 직접 쓴 루프의 결과가 같은가:",
      all(torch.equal(p, q) for p, q in zip(a.parameters(), b.parameters())))

trainer와 직접 쓴 루프의 결과가 같은가: True

비트 단위로 같다. 이 비교는 trainer를 고칠 때마다 다시 돌리는 회귀 테스트(6.6장)로 남겨 둔다.

확장 지점

trainer를 재사용하다 보면 "이 실험에서만 필요한 기능"이 계속 생긴다. 매번 trainer에 if 플래그를 추가하면 trainer는 곧 모든 실험의 예외를 담은 거대한 함수가 된다. 확장 방법을 미리 정해 두면 이를 막을 수 있다. 확장 지점은 세 종류다.

• step 함수 교체: 과제마다 forward와 loss가 다를 때. 분류·회귀·자기지도 학습·물리 손실 항 추가(24.2장)가 여기에 속한다.
• callback: 루프의 흐름은 그대로 두고 특정 시점에 무언가를 관찰하거나 저장하거나 멈추게 할 때. 로깅, checkpoint, early stopping, 학습률 기록, 예측 샘플 저장, 일정 epoch 후 일부 층 동결 해제.
• trainer 상속이나 새 trainer: 루프의 구조 자체가 다를 때. 한 step에 optimizer가 둘인 학습(GAN), 안쪽·바깥쪽 루프가 있는 이중수준 최적화(10.7장), 환경과 상호작용하는 강화학습(20장)은 callback으로 흉내 내지 말고 별도 루프를 쓴다.

callback의 이벤트는 많을 필요가 없다. 실습의 trainer.py는 on_train_start, on_step_end, on_eval_end, on_train_end 네 개로 로깅·저장·조기 종료를 모두 처리한다. 이벤트를 추가하기 전에 기존 이벤트로 할 수 없는지 먼저 묻는다.

callback 설계에서 지킬 규칙이 둘 있다.

• callback끼리 서로를 몰라야 한다. checkpoint callback이 early stopping callback의 내부 값을 읽기 시작하면 순서 의존성이 생긴다. 공유가 필요한 값은 trainer의 상태(예: last_metrics)를 통해 주고받는다.
• 상태가 있는 callback은 state_dict를 가진다. early stopping의 best와 기다린 횟수는 학습 상태의 일부다. 재개할 때 이것을 복원하지 않으면 재개 직후의 첫 검증이 무조건 "개선"으로 판정되고 patience가 처음부터 다시 센다. 실습의 trainer는 callback의 상태를 checkpoint에 함께 저장한다.

eval.py의 분리

평가 코드는 학습 코드 없이 실행할 수 있어야 한다. 이유는 연구의 시간 순서에 있다.

• 재평가. metric 계산의 버그를 고치거나 새 metric을 추가하면, 이미 학습한 checkpoint를 다시 학습하지 않고 모두 다시 평가해야 한다.
• 공정한 비교. baseline, ablation, 다른 시드의 checkpoint를 같은 평가 코드로 평가해야 표의 숫자들이 비교 가능하다(6.1장).
• 새 데이터. 몇 달 뒤 새로 측정한 데이터에 같은 모델을 적용할 때, 그때의 학습 코드는 이미 바뀌어 있을 수 있다.

그래서 eval.py가 import하는 것은 모델 정의, 데이터 읽기, metric, checkpoint 읽기뿐이어야 한다. trainer, optimizer, 학습 설정의 기본값에 의존하지 않는다. 학습 중에 계산한 검증 metric과 eval.py가 다시 계산한 값이 같은지 확인하는 것이 이 분리가 제대로 되었는지에 대한 시험이다. 실습의 eval.py는 일부러 학습 때와 다른 batch 크기(7)로 평가해, 집계가 batch 크기에 의존하지 않는지도 함께 확인한다.

프레임워크

직접 쓴 trainer가 맡는 일을 대신 해 주는 학습 프레임워크가 여럿 있다. 차이는 얼마나 많은 것을 가져가는가, 즉 추상화 수준이다. 아래 설명은 각 도구의 역할을 일반적으로 정리한 것이다(집필 시점 예시. 이 책의 실습 환경에는 설치되어 있지 않으며, 구체적인 API와 기본값은 각 도구의 공식 문서에서 확인한다).

• Accelerate는 루프를 사용자가 그대로 쓰게 두고, 같은 코드를 CPU·GPU 하나·여러 GPU에서 돌게 만드는 얇은 층이다. 직접 쓴 trainer와 함께 쓸 수 있다.
• Lightning은 이 장의 구조(trainer가 루프를 갖고, 사용자는 step을 정의하고, 부가 기능은 callback)와 같은 생각을 더 넓게 구현한 프레임워크다. 그래서 직접 trainer를 설계해 본 사람은 Lightning의 구조를 빨리 읽는다.
• **Transformers의 Trainer**는 그 라이브러리의 모델을 학습하는 데 맞춰져 있다. 그 모델 계열에서는 코드가 가장 짧지만, 루프 구조가 다른 연구에는 맞지 않는다.

어느 도구든 기본값이 많다. 검증 주기, 로깅 주기, gradient clipping 여부, 정밀도, checkpoint 저장 기준, 학습률 스케줄과 warm-up, 데이터 셔플 방식이 모두 기본값을 갖는다. 기본값은 대부분 합리적이지만, 그것이 내 실험의 설정이 된다는 사실은 변하지 않는다. PyTorch 자체에도 이런 기본값이 있다. optimizer를 Adam에서 AdamW로 바꾸면서 인자를 그대로 두면 무엇이 달라지는가.

for opt_cls in (torch.optim.Adam, torch.optim.AdamW):
    model, _, data = setup()
    opt = opt_cls(model.parameters(), lr=1e-2)                 # 다른 인자는 주지 않았다
    fit(model, opt, data, regression_step, max_epochs=200)
    norm = torch.cat([p.detach().flatten() for p in model.parameters()]).norm()
    loss = sum(regression_step(model, batch).item() for batch in data) / len(data)
    print(f"{opt_cls.__name__:5s} weight_decay 기본값 {opt.defaults['weight_decay']:<5}  "
          f"파라미터 노름 {norm:.4f}  학습 손실 {loss:.4f}")

Adam  weight_decay 기본값 0      파라미터 노름 6.8874  학습 손실 0.0120
AdamW weight_decay 기본값 0.01   파라미터 노름 6.7403  학습 손실 0.0135

AdamW는 weight decay를 지정하지 않아도 0.01을 적용한다(집필 시점 예시. PyTorch 2.15에서 확인). "optimizer만 바꿨다"고 기록한 실험이 실제로는 regularization도 함께 바꾼 실험이 된다. 이 작은 문제에서는 차이가 작지만, 학습률이 크거나 학습이 길면 커진다(10.4장). 프레임워크의 기본값은 이런 항목이 수십 개다.

도입 판단

프레임워크를 쓸지는 취향이 아니라 비용의 비교다. 프레임워크는 반복 코드를 줄여 주는 대신 숨겨진 동작을 늘린다. 판단 기준은 다음과 같다.

도입하면 이득이 큰 경우

• 여러 GPU·여러 노드, mixed precision, gradient accumulation처럼 표준화되어 있고 틀리기 쉬운 반복 코드가 필요할 때. 직접 쓴 분산 학습 코드의 버그는 찾기 어렵다(22.6장).
• 학습 루프 자체는 연구 대상이 아니고, 모델이나 데이터가 연구 대상일 때.
• 팀이 이미 그 도구로 일하고 있고, 코드를 서로 읽어야 할 때.

직접 쓰는 편이 나은 경우

• 루프의 구조가 연구 대상일 때: 새 optimizer, 이중수준 최적화, 비표준 업데이트 순서, step마다 다른 데이터 소스를 섞는 학습. 프레임워크의 추상화를 우회하는 코드가 프레임워크 없이 쓴 코드보다 길어진다.
• 학습 루프를 처음 배우는 단계. 프레임워크가 대신하는 일을 모르면 프레임워크가 실패할 때 할 수 있는 것이 없다. 이 부가 루프를 직접 쓰는 것부터 시작한 이유다.
• 실험이 CPU나 GPU 하나에서 끝나는 작은 규모일 때. 이때 반복 코드는 100–200줄이고, 그 줄은 모두 내가 이해한다.

도입한다면 지킬 것

• 사용하는 버전을 고정하고(3.3장), 버전을 올릴 때 결과가 바뀌는지 확인한다.
• 해석이 끝난 설정 전체, 즉 내가 지정하지 않은 기본값까지 포함한 설정을 실행마다 기록한다(6.3장). 대부분의 프레임워크는 이를 출력하거나 저장하는 방법을 제공한다.
• 직접 쓴 짧은 루프와 같은 결과를 내는지, 작은 문제에서 한 번 비교한다. 다르면 차이를 설명할 수 있을 때까지 넘어가지 않는다.
• 프레임워크 밖에서 돌아가는 eval.py를 유지한다. 학습은 프레임워크로 해도 평가는 프레임워크 없이 checkpoint만으로 재현할 수 있어야 한다.