이제 지금까지 만난 네 가지 변환인 jit, grad, vmap, pmap을 바탕으로 첫 번째 완결된 프로그램을 만들어 보자. 잡음이 섞인 데이터에 다항식을 맞추는 작은 학습기가 목표야.
import jax
import jax.numpy as jnp
from jax import random
# 1. 합성 데이터 만들기
key = random.PRNGKey(42)
key, x_key, noise_key = random.split(key, 3)
x = random.uniform(x_key, (100,), minval=-5, maxval=5)
true_a, true_b, true_c = 2.0, -1.5, 3.0
y = true_a * x**2 + true_b * x + true_c + 0.5 * random.normal(noise_key, (100,))
# 2. model — 그냥 함수
def model(params, x):
a, b, c = params
return a * x**2 + b * x + c
# 3. loss — 또 함수
def loss(params, x, y):
pred = model(params, x)
return jnp.mean((pred - y) ** 2)
# 4. gradient + jit — 하나로 묶음
grad_fn = jax.jit(jax.grad(loss))
# 5. 학습 loop
params = jnp.array([0.0, 0.0, 0.0])
lr = 0.01
for step in range(2000):
g = grad_fn(params, x, y)
params = params - lr * g
print(f"최종 params: {params}")
print(f"정답: [{true_a}, {true_b}, {true_c}]")
🌱 퀘스트의 출발점
이 30줄짜리 학습기의 구조를 이해하면 뒤의 13개 트랙을 따라갈 기반이 생겨. 신경망도 결국 더 큰 함수일 뿐이야. 매개변수가 커지고 모델이 복잡해지고 옵티마이저 라이브러리가 더해져도, 모두 이 틀을 확장한 것이지 완전히 다른 패러다임은 아니야.
Code
import jax
import jax.numpy as jnp
# 1. Generate synthetic data
key = jax.random.PRNGKey(42)
key, subkey = jax.random.split(key)
X = jax.random.normal(subkey, shape=(100, 3)) # 100 samples, 3 features
true_w = jnp.array([2.0, -1.0, 0.5])
y = X @ true_w + 0.1 * jax.random.normal(key, shape=(100,))
# 2. Define the model and loss as pure functions
def predict(params, x):
return jnp.dot(x, params)
def loss_fn(params, X, y):
preds = predict(params, X)
return jnp.mean((preds - y) ** 2)
# 3. Compile the gradient computation
@jax.jit
def update(params, X, y, lr=0.1):
grads = jax.grad(loss_fn)(params, X, y)
return params - lr * grads # Simple gradient descent
# 4. Train
params = jnp.zeros(3) # Start from zeros
for step in range(100):
params = update(params, X, y)
if step % 20 == 0:
current_loss = loss_fn(params, X, y)
print(f"Step {step}, Loss: {current_loss:.4f}")
print(f"Learned params: {params}")
print(f"True params: {true_w}")
# Per-example gradients: gradient of loss for EACH data point
# In PyTorch, this requires special tricks. In JAX, it's one line.
def single_loss(params, x, y):
"""Loss for a single example."""
pred = jnp.dot(x, params)
return (pred - y) ** 2
# vmap over the data dimensions (axis 0 of x and y), not params
per_example_grad_fn = jax.vmap(jax.grad(single_loss), in_axes=(None, 0, 0))
per_example_grads = per_example_grad_fn(params, X, y)
print(per_example_grads.shape) # (100, 3) — one gradient per example
처음부터 스크립트를 작성해 y = ax² + bx + c를 잡음이 섞인 100개 점에 맞춰 봐. jax.grad와 수동 그래디언트 하강으로 1000단계 학습하고, jit으로 컴파일한 경우와 eager 실행의 단계별 시간을 비교해. 수렴한 뒤 학습한 계수와 정답 계수를 함께 출력해.
Progress
Progress is local-only — sign in to sync across devices.