PyTorch Lightning, Keras의 model.fit(...) 같은 상위 수준 래퍼가 JAX 코어에는 없어. 모든 학습 스텝을 직접 작성해. 처음에는 답답하지만 의도된 설계야.
JAX 식 학습 루프의 모양
@jax.jit
def train_step(state, batch):
x, y = batch
def loss_fn(params):
pred = model.apply(params, x)
loss = compute_loss(pred, y)
metrics = {"acc": accuracy(pred, y)}
return loss, metrics
(loss, metrics), grads = jax.value_and_grad(loss_fn, has_aux=True)(state.params)
updates, new_opt_state = optimizer.update(grads, state.opt_state, state.params)
new_params = optax.apply_updates(state.params, updates)
new_state = state.replace(
params=new_params,
opt_state=new_opt_state,
step=state.step + 1,
)
return new_state, loss, metrics
# 사용자가 직접 loop
for batch in dataloader:
state, loss, metrics = train_step(state, batch)
if state.step % 100 == 0:
print(f"step {state.step}: loss={loss:.4f}")
30줄이면 충분하고 모든 과정이 보여. 마법처럼 숨은 곳도 없어.
왜 이게 좋은가?
- 모든 step이 가시적: 그래디언트, 옵티마이저 상태, 매개변수 갱신을 모두 직접 확인할 수 있어. 학습기.fit() 안에 숨어 있던 부분이 모두 드러나.
- 커스터마이징 자유: 그래디언트 자르기, custom 갱신 규칙, 다양한 스케줄, 모두 같은 30줄 안에 추가해. 래퍼 API의 한계 없어.
- 디버깅 단순: print, assert, breakpoint, 어디든 자유. 마법 같은 callback 시스템 없어.
- JIT 명확: train_step 함수가 정확히 무엇을 jit하는지 드러나고 컴파일 비용도 확인할 수 있어.
비교, PyTorch Lightning
# PyTorch Lightning
class MyModel(pl.LightningModule):
def training_step(self, batch, batch_idx):
x, y = batch
pred = self.model(x)
loss = F.cross_entropy(pred, y)
self.log("train_loss", loss)
return loss
def configure_optimizers(self):
return torch.optim.AdamW(self.parameters(), lr=1e-3)
trainer = pl.Trainer(max_epochs=10, gpus=4)
trainer.fit(model, dataloader)
편리하지만 무엇이 어떻게 돌아가는지 알려면 Lightning의 소스 코드를 살펴봐야 해. JAX에서는 직접 쓴 30줄이 곧 소스 코드야.
고수준 래퍼가 필요할 때
Flax의 nnx.training, chex, clu 같은 라이브러리가 일부 상용구를 줄여 줘. 그러나 핵심은 같은 명시적 패턴이고 보조 도구만 추가해. PyTorch Lightning 같은 통합 래퍼가 JAX 코어에 없는 건 의도된 설계야.
🛠 JAX 학습 코드의 mantra
"Show me the 루프." JAX 코드를 받으면 train_step 함수 한 개 + 루프 한 개. 30줄로 끝이야. 래퍼가 늘어날수록, 그 코드가 JAX 답지 않는지 의심해 봐. 학습 코드가 짧아지는 것보다, 명료한 게 더 가치 있다는 게 JAX 공동체의 가치관.
한 가지, 명시적 학습 코드는 처음 익히는 데 학습 곡선이 있어. 몇 번 짜 보면 패턴이 눈에 들어와서, 새 작업마다 빠르게 응용할 수 있어. PyTorch Lightning의 콜백들을 외우는 것보다, 다른 작업에도 더 잘 옮겨 갈 수 있는 지식이야.