Training Loop 해부

모든 진지한 training loop 가 가진 6 가지

Config — 모든 hyperparameter capture 하는 dataclass 또는 YAML/TOML. 모든 checkpoint 에 저장.
Seed — 맨 위 set, 모든 거에 적용.
Mixed precision — autocast + GradScaler (또는 BF16, scaler 없음).
Gradient clipping — backward 와 optimizer step 사이.
Learning-rate schedule — transformer 에 warmup + cosine, fine-tuning 에 OneCycle, old-school CNN 에 step decay.
Validation + best-checkpoint saving — 매 epoch (또는 N step) 마다, held-out set 에 eval, val metric 개선되면 저장.

팁: 이 6 가지 이미 포함한 'blank training loop' template 부터 모든 project 시작하면, 모든 project 마다 같은 디버깅 한 주 절약.

Gradient accumulation pattern

VRAM 에 원하는 batch size 가 너무 큰 model: 여러 forward/backward 동안 opt.zero_grad() 건너뛰고 step + zero. Single 더 큰 batch 와 mathematically equivalent. Loss magnitude consistent 유지하려 scaler.scale(loss / accum_steps).backward() 조정.

Framework 가 도움 되는 곳

PyTorch Lightning, Hugging Face Trainer, Accelerate 가 boilerplate (mixed precision, distributed, checkpointing, logging) 처리해서 본인이 model 과 data 에 집중 가능. Trade-off: 적은 코드, 적은 visibility. Framework 채택 전 적어도 한 번 본인 loop 손으로 작성.

원칙: 첫 training loop 는 손으로. 10 번째는 아마 framework. 중간이 가장 많이 배우는 곳.

Code

A complete production-quality training loop·python

import torch, json
from torch import nn, optim
from torch.cuda.amp import autocast, GradScaler
from torch.nn.utils import clip_grad_norm_

def train(model, train_loader, val_loader, cfg, device):
    model = model.to(device)
    opt = optim.AdamW(model.parameters(), lr=cfg.lr, weight_decay=cfg.wd)
    sch = warmup_cosine(opt, cfg.warmup_steps, cfg.total_steps)
    scaler = GradScaler()
    loss_fn = nn.CrossEntropyLoss()

    best_val = 0.0
    step = 0
    for epoch in range(cfg.epochs):
        model.train()
        for xb, yb in train_loader:
            xb, yb = xb.to(device), yb.to(device)
            opt.zero_grad()
            with autocast(dtype=torch.bfloat16):
                logits = model(xb)
                loss = loss_fn(logits, yb)
            scaler.scale(loss).backward()
            scaler.unscale_(opt)
            clip_grad_norm_(model.parameters(), max_norm=cfg.clip)
            scaler.step(opt); scaler.update(); sch.step()
            step += 1

        val_acc = evaluate(model, val_loader, device)
        if val_acc > best_val:
            best_val = val_acc
            torch.save({"model": model.state_dict(), "config": cfg.__dict__,
                        "val_acc": val_acc, "step": step}, "best.pt")
        print(f"epoch={epoch} step={step} val={val_acc:.4f}")
    return best_val

모든 진지한 training loop 가 가진 6 가지

Gradient accumulation pattern

Framework 가 도움 되는 곳

Code

External links

Exercise

Progress

댓글 0