Skip to content

Objectives and metrics

A model may learn several physical quantities whose errors need different averaging weights. The engine keeps those choices explicit: a step returns the loss and sample count, while metrics describe the statistics to report. Loss normalization follows the data contributing to an optimizer update, including batches of different sizes and distributed execution.

A mean loss

loss is a local mean. Its denominator defaults to num_samples; pass loss_weight when it is an atom count, number of labeled components, or another explicit weight. Each update divides the weighted loss sum by the actual global denominator across its accumulation window, including incomplete tails. Terms combined into one scalar mean must share the declared denominator. Use Objective for independently normalized terms, as shown below.

The training example uses this form. StepOutput also carries metrics for training or evaluation, independently of the loss used for backpropagation.

StepOutput dataclass

StepOutput(
    loss: Tensor | Mapping[str, Tensor],
    num_samples: int,
    metrics: dict[str, Metric] = dict(),
    loss_weight: Scalar | None = None,
)

Output of one task-defined training or evaluation step.

Parameters:

  • loss (Tensor | Mapping[str, Tensor]) –

    A mean over loss_weight units, or named differentiable sums when the Trainer/Evaluator has an Objective. Training retains its graph.

  • num_samples (int) –

    Number of samples in the local batch.

  • metrics (dict[str, Metric], default: dict() ) –

    Additional scalar statistics; loss names and prefixes are reserved.

  • loss_weight (Scalar | None, default: None ) –

    Denominator of the local mean objective, defaulting to num_samples. Use e.g. atom count for an atom-averaged objective. Scalar composite terms must share this denominator. Ignored with Objective, which supplies independently counted named terms.

Independently normalized objectives

Here the model returns two predictions and the loaders yield dictionaries containing their targets. The accelerator and optimizer are configured as in Training and optimization. Each output has its own normalization count; batch["mask"] is a boolean mask with the same shape as the auxiliary targets, identifying their available components.

from elfes.engine import Evaluator, Objective, StepOutput, Trainer

objective = Objective(
    weights={"primary": 1.0, "auxiliary": 0.2},
    counts=lambda batch: {
        "primary": batch["target"].numel(),
        "auxiliary": batch["mask"].sum(),
    },
)


def step(model, batch):
    primary, auxiliary = model(batch["inputs"])
    mask = batch["mask"]
    return StepOutput(
        loss={
            "primary": (primary - batch["target"]).square().sum(),
            "auxiliary": (auxiliary[mask] - batch["aux_target"][mask]).square().sum(),
        },
        num_samples=len(batch["inputs"]),
    )


evaluator = Evaluator(accelerator, {"val": val_loader}, step, objective=objective)
trainer = Trainer(
    accelerator, model, optimizer, train_loader, step,
    objective=objective, evaluator=evaluator,
)

With an Objective, loss values are differentiable sums, and coefficients live in objective.weights. The engine sums each term's counts over the actual global accumulation window, then optimizes the weighted sum of term means. counts(batch) must read deterministic, nonnegative, finite data counts without model execution, random sampling, or mutation. Counts may be fractional but cannot depend on trainable parameters.

This path buffers one window of input batches on the execution device, while retaining only one microbatch's forward graph. The scalar mean path streams batches directly. Mask out missing targets before forming the loss; an empty indexed error sum supplies a differentiable zero. Globally absent terms contribute no gradient. A wholly unlabelled window skips optimizer, EMA, and step scheduler, including momentum and weight decay. Evaluation reports zero counts and omits unobserved means instead of reporting perfect zero error.

Metrics include loss/<name>, loss_count/<name>, and the weighted total loss. Loss names and the loss_sum/ internal prefix are reserved. Additional physical metrics remain independent of the training objective.

Objective dataclass

Objective(weights: Mapping[str, float], counts: Callable[[Any], Mapping[str, Scalar]])

Optimize a weighted sum of independently normalized loss terms.

counts(batch) returns each term's nonnegative normalization count, without a model forward or RNG use. StepOutput.loss returns matching differentiable sums, including a differentiable zero when no labels are present. Weights and counts must not depend on trainable parameters.

Reporting metrics

Metric updates carry sufficient statistics so the engine can combine batches and ranks before computing a reported value. For a mean, supply the sum and its count; for a root mean square, supply the sum of squared values and its count. Sums and maxima accumulate their corresponding quantities directly. This preserves the intended weighting when batch sizes differ.

For example, MeanMetric(error.abs().sum(), error.numel()) reports a component-averaged absolute error. The count can differ from the number of samples and from the denominator used by the training loss. The evaluation reference describes how these statistics cover uneven or empty distributed shards.

MeanMetric dataclass

MeanMetric(total: Scalar, count: Scalar)

Sufficient statistics for a weighted mean.

RootMeanSquareMetric dataclass

RootMeanSquareMetric(sum_squared: Scalar, count: Scalar)

Sufficient statistics for a root mean square.

SumMetric dataclass

SumMetric(total: Scalar)

Sufficient statistics for a sum.

MaxMetric dataclass

MaxMetric(value: Scalar)

Sufficient statistics for a maximum.