Flax 소개 — JAX용 신경망 라이브러리

Flax 소개

Flax는 "Neural Networks for JAX"의 줄임말이에요. JAX를 사용하는 연구자와 개발자에게 유연한 end-to-end 경험을 제공하는 신경망 라이브러리예요. JAX의 전체 기능을 그대로 활용할 수 있게 해주면서, 네트워크를 더 쉽게 표현·디버깅할 수 있게 도와줘요.

출처: https://flax.readthedocs.io/en/latest/

NNX 기본 사용

Flax NNX는 파라미터 같은 상태를 모듈이 직접 들고 다니는 방식이에요. 학습 스텝은 옵티마이저와 nnx.jit 변환으로 간단히 표현돼요.

optimizer = nnx.Optimizer(model, optax.adam(1e-3), wrt=nnx.Param)

@nnx.jit  # automatic state propagation
def train_step(model, optimizer, x, y):
    loss_fn = lambda model: ((model(x) - y) ** 2).mean()
    loss, grads = nnx.value_and_grad(loss_fn)(model)

설치

pip install flax

최신 저장소 버전은 git에서 직접 설치할 수 있어요.

pip install git+https://github.com/google/flax.git

배우기

• Flax NNX Basics — nnx_basics.html
• MNIST Tutorial — mnist_tutorial.html
• Guides — <guides/index.html>
• Flax Linen to Flax NNX — guides/linen_to_nnx.html
• API reference — api_reference/index.html

더 알아보기

• Flax 문서
• Flax NNX Basics
• Flax GitHub