Training and evaluation#

Write an explicit optimization step, clip gradients, separate training from evaluation, and aggregate held-out metrics.

These outputs were saved from local execution. The documentation renderer does not execute notebook cells.

Hide code cell source

 
from pathlib import Path
import sys
ROOT = Path.cwd() if (Path.cwd() / 'examples').exists() else Path.cwd().parent
sys.path.insert(0, str(ROOT))
import jax
import jax.numpy as jnp
import numpy as np
import jaxsaw as js
from jaxsaw import nn, optim, data, checkpoint, inspection
from jaxsaw.nn import functional as F
print(f"JAXsaw {js.__version__} · JAX {jax.__version__}")
print(jax.devices())

JAXsaw 0.1.0 · JAX 0.5.3
[CudaDevice(id=0)]

Prepare a controlled classification task#

The two-dimensional synthetic task makes the loop easy to inspect. Training and held-out examples use independent random draws. This is an API exercise, not a dataset benchmark.

rng = np.random.default_rng(8)
x = rng.normal(size=(256,2)).astype(np.float32)
y = (x[:,0] + .5*x[:,1] > 0).astype(np.int32)
train = (jnp.asarray(x[:192]), jnp.asarray(y[:192]))
valid = (jnp.asarray(x[192:]), jnp.asarray(y[192:]))
class Classifier(nn.Module):
    def __init__(self, *, rngs):
        super().__init__()
        self.hidden = nn.Linear(2,16,rngs=rngs)
        self.dropout = nn.Dropout(.1)
        self.head = nn.Linear(16,2,rngs=rngs)
    def forward(self,x,*,rngs=None):
        return self.head(self.dropout(F.relu(self.hidden(x)),rngs=rngs))
model = Classifier(rngs=js.Rngs(params=8))
optimizer = optim.Adam(model.parameters(),lr=.02)
runtime = js.Rngs(dropout=9)

Differentiate, clip, update#

The step returns gradients explicitly, then applies one optimizer update. Auxiliary predictions supply metrics without a second stochastic forward pass.

@js.jit
def train_step(model, optimizer, batch, rngs):
    x,y = batch
    def objective(candidate):
        logits = candidate(x,rngs=rngs)
        return F.cross_entropy(logits,y), (logits.argmax(-1)==y).sum()
    (loss,correct), grads = js.value_and_grad(objective,has_aux=True)(model)
    grads,norm = optim.clip_grad_norm(grads,1.)
    optimizer.step(grads)
    return {'loss':loss,'correct':correct,'count':y.size,'gradient_norm':norm}
model.train()
for step in range(100):
    result = train_step(model,optimizer,train,runtime)
print({name:float(value) for name,value in result.items()})
print('Optimizer updates:',int(optimizer.step_count))

{'correct': 191.0, 'count': 192.0, 'gradient_norm': 0.02445269748568535, 'loss': 0.02627979777753353}
Optimizer updates: 100

Evaluate without updating state#

Evaluation mode changes dropout and normalization behavior; it is not a gradient switch. Validation uses an explicit, separate function and a held-out batch.

@js.jit
def evaluate(model,batch):
    x,y = batch
    logits = model(x)
    return {'loss':F.cross_entropy(logits,y),
            'correct':(logits.argmax(-1)==y).sum(),'count':y.size}
model.eval()
result = evaluate(model,valid)
print('Held-out loss:',float(result['loss']))
print('Held-out accuracy:',float(result['correct']/result['count']))

Held-out loss: 0.031200522556900978
Held-out accuracy: 1.0

Aggregate by valid item count#

Do not average batch means with equal weight when the final batch is smaller. Masked sequence tasks should report valid-token counts.

from jaxsaw.metrics import ClassificationAccumulator
accumulator = ClassificationAccumulator()
for start in range(0,len(valid[0]),15):
    batch = tuple(array[start:start+15] for array in valid)
    accumulator.update(evaluate(model,batch))
print(accumulator.compute())

{'loss': 0.031184522405965254, 'accuracy': 1.0}

Adapt the objective#

Change the scalar objective for regression, reconstruction, or multiple supervised heads. In adversarial training, use separate optimizers and select the differentiated module explicitly. See the Fashion-MNIST example for both phases.