Metrics API

torch_batteries.utils.metrics

Metric protocol, implementations, and helper exports.

PhaseMetricManager Source

Coordinate callable, incremental, and collected metrics for one phase.

Ordinary callables produce batch values that are sample-weighted by progress tracking. Stateful metrics own their exact aggregation. CollectedMetric instances share detached CPU collections. Metric lifecycle failures either propagate immediately or skip the failed metric for the remainder of the phase.

Parameters:

Name	Type	Description	Default
metrics	dict[str, Metric]	Named callable, stateful, or collected metrics.	required
metric_error_policy	Literal['raise', 'warn']	"raise" to propagate metric exceptions or "warn" to log and skip a failed metric for the current phase.	'raise'

reset() Source

Reset all phase-scoped metric state.

update(predictions, targets) Source

Update phase metrics and return per-batch callable values.

Parameters:

Name	Type	Description	Default
predictions	Tensor	Model predictions for the batch.	required
targets	Tensor	Targets for the batch.	required

compute() Source

Compute all full-phase metric values.

state_dict() Source

Return optional states exposed by configured metric objects.

load_state_dict(state_dict) Source

Restore optional configured metric states strictly by name.

Parameters:

Name	Type	Description	Default
state_dict	dict[str, Any]	Serialized state keyed by configured metric name.	required

CollectedMetric Source

Bases: StatefulMetric

Evaluate an ordinary metric callable once over a complete phase.

This convenience adapter retains detached CPU predictions and targets, so its memory use grows with the dataset. Prefer an incremental StatefulMetric implementation for large datasets.

Parameters:

Name	Type	Description	Default
metric	MetricCallable	Callable evaluated once with concatenated phase tensors.	required

reset() Source

Clear retained phase tensors.

update(predictions, targets) Source

Retain one detached CPU batch.

Parameters:

Name	Type	Description	Default
predictions	Tensor	Model predictions for the batch.	required
targets	Tensor	Targets for the batch.	required

compute() Source

Concatenate retained tensors and evaluate the wrapped callable.

compute_collected(predictions, targets) Source

Evaluate the wrapped callable over shared collected tensors.

Parameters:

Name	Type	Description	Default
predictions	Tensor	Concatenated phase predictions.	required
targets	Tensor	Concatenated phase targets.	required

StatefulMetric Source

Bases: Protocol

Protocol for metrics computed from state accumulated over a full phase.

Battery calls reset before each phase, update with detached predictions and targets for each batch, then compute once. Implementations should return one numeric scalar.

reset() Source

Reset metric state before a phase.

update(predictions, targets) Source

Update metric state with one batch.

Parameters:

Name	Type	Description	Default
predictions	Tensor	Detached model predictions for the batch.	required
targets	Tensor	Detached targets for the batch.	required

compute() Source

Compute a scalar metric from accumulated state.

calculate_metrics(metrics, pred, target) Source

Calculate multiple metrics for given predictions and targets.

This function takes a dictionary of metric functions and applies them to the predictions and targets. Each metric function should accept two tensors (predictions and targets) and return a scalar value (either as a tensor or float).

The function handles both tensor and scalar returns from metric functions, automatically converting tensors to Python floats using .item().

If a metric function raises an exception during calculation, the error is logged as a warning and the metric is skipped (not included in the returned dictionary).

Parameters:

Name	Type	Description	Default
metrics	dict[str, MetricCallable]	Dictionary mapping metric names to callable functions. Each function should have signature: fn(pred, target) -> float | Tensor	required
pred	Tensor	Model predictions as a tensor	required
target	Tensor	Ground truth target values as a tensor	required

Returns:

Type	Description
dict[str, float]	Dictionary mapping metric names to their calculated float values.
dict[str, float]	Only successfully calculated metrics are included.

Examples:

 import torch.nn.functional as F

 def mae(pred, target):
     return F.l1_loss(pred, target)

 def rmse(pred, target):
     return torch.sqrt(F.mse_loss(pred, target))

 metrics_dict = {'mae': mae, 'rmse': rmse}
 pred = torch.tensor([[1.0], [2.0], [3.0]])
 target = torch.tensor([[1.1], [2.2], [2.9]])

 results = calculate_metrics(metrics_dict, pred, target)
 # returns: {'mae': 0.133..., 'rmse': 0.141...}

Note

• Metric functions should not modify the input tensors
• Both pred and target should have compatible shapes for the metric functions
• Failed metric calculations are logged but don't raise exceptions