Need readable training loops without giving up research flexibility? Common Loop Utils (CLU) provides JAX- and Flax-compatible utilities, especially functional metric computation. section 11.1.3 section 12.1.3
CLU metrics follow an accumulate-then-compute pattern:
- Compute local values from each batch’s model outputs, such as loss, logits, and labels. section 11.1.3
- Connect metrics to named outputs with interfaces such as Metric.from_output(). section 11.1.3
- Aggregate intermediate values across batches with merge(), then obtain final values with compute(). section 11.1.3
- Store a metric collection alongside Flax TrainState so updates remain part of the training state. section 11.1.3
Adding metrics to training state
Define a Metrics collection with metrics.Accuracy and metrics.Average.from_output('loss'), initialize it with Metrics.empty(), and update it using state.metrics.single_from_model_output(...), merge(), and state.replace(metrics=metrics). Convert labels to int32 when required by CLU. section 11.1.3
Listing 11.6 Adding metrics to the training state
from flax.training import train_state
from clu import metrics
import flax
import optax
@flax.struct.dataclass #1
class Metrics(metrics.Collection): #1
accuracy: metrics.Accuracy #2
loss: metrics.Average.from_output('loss') #3
class TrainState(train_state.TrainState):
metrics: Metrics #4
state = TrainState.create(
apply_fn=model.apply,
params=params,
tx=optax.sgd(learning_rate=0.01, momentum=0.9),
metrics=Metrics.empty()) #5
#1 Defines our own class for metrics
#2 Adds an accuracy metric
#3 Adds a loss metric
#4 Adds metrics to the TrainState
#5 Initializes metrics with an empty collection
Inspect the compact pattern for defining, attaching, and initializing CLU metrics in a Flax training state. section 11.1.3
