Skip to content

Evaluation

Evaluation measures a model over named data splits and supplies the metrics used to select checkpoints or stop training. It can run periodically inside training or independently with a loaded model. The engine counts every real sample once and restores the training model's mode and random state afterwards.

Raw and averaged parameters

The trainer in the training example maintains an exponential moving average (EMA) of its parameters. Its current raw or averaged parameters can be evaluated explicitly:

raw_metrics = trainer.evaluate(weights="raw")
ema_metrics = trainer.evaluate(weights="ema")

Periodic training evaluation uses EMA when enabled, and best-model selection follows those metrics. Explicit evaluation and export APIs default to raw weights. An EMA request fails when EMA is absent; it never silently falls back to raw.

Evaluator requires drop_last=False and counts every real sample once, including uneven or empty rank shards. It restores model mode and training RNG after evaluation. Set requires_grad=True for predictions derived through autograd, such as conservative forces. Distributed evaluation unwraps DDP for independent forwards while retaining compilation and mixed precision; evaluation steps must not perform their own distributed collectives.

For a saved model, load checkpoint weights and call evaluator.evaluate(model) directly. Standalone evaluation needs no optimizer or training DataLoader.

Evaluator 

Evaluator(
    accelerator: Accelerator,
    dataloaders: Mapping[str, DataLoader[Any]],
    step: Step,
    *,
    requires_grad: bool = False,
    objective: Objective | None = None,
)

Evaluate every sample once, including uneven and empty rank shards.

DDP forwards are unwrapped; model state is broadcast once before evaluation. Steps may use input autograd but must not perform distributed collectives. Dataset reads and collate functions must be deterministic for exact resume.

Best-model selection and early stopping

Monitor identifies a metric by its split and name, such as validation/loss. Its mode chooses whether smaller or larger values are better, and min_delta sets the required improvement. An optional patience stops training after that many evaluations without improvement. The selected metric must be finite.

Pass the monitor through TrainConfig(monitor=...). When a checkpoint directory is configured, an improvement updates best; the checkpoint reference explains how to select raw or EMA parameters from that position. For a ReduceLROnPlateau scheduler, the trainer uses the same monitored metric.

Monitor dataclass 

Monitor(
    metric: str,
    mode: Literal["min", "max"] = "min",
    patience: int | None = None,
    min_delta: float = 0.0,
)

Best-model and optional early-stopping policy.