xnn.common.benchmark.metrics.collect_predictions#

xnn.common.benchmark.metrics.collect_predictions(model, loader, device, targets, atomic_energies=None, energy_per_atom=True)[source]#

Run model over loader once, pairing predictions with targets.

The model is put in eval mode but gradients are left enabled, because force predictions differentiate the energy with respect to positions (the same reason Trainer.evaluate() avoids torch.no_grad). Each prediction/target pair is detached and moved to CPU before being stacked.

Energy handling is controlled by atomic_energies and energy_per_atom. When atomic_energies is given, the per-element reference energy of every atom is subtracted from both the predicted and reference total energy, turning them into atomization (interaction) energies – the physically meaningful quantity to report (see energy). When energy_per_atom is set, the energy is then divided by the per-structure atom count, the size-extensive normalization the training loss uses, so energy metrics are comparable across differently sized structures. Both are applied identically to the prediction and the reference, so difference metrics (MAE/MSE/RMSE) are invariant to the E0 offset while their reported values become meaningful.

Parameters:

model (torch.nn.Module) – A model whose forward returns a dict with an "energy" key and, when forces/stress are requested, "forces" / "stress" keys (e.g. a ForceStressOutput).
loader (torch.utils.data.DataLoader) – Loader yielding batched AtomicGraph objects.
device (torch.device) – Device to evaluate on.
targets (list of str) – Which quantities to collect; a subset of "energy", "forces" and "stress". A target is skipped for a batch that lacks the reference value or the corresponding prediction.
atomic_energies (torch.Tensor or None, optional) – A Z-indexed lookup of per-element reference energies (as built by build_e0_lookup()). When given, energies are scored on an atomization basis. Defaults to None (raw total energy).
energy_per_atom (bool, optional) – Whether to divide the (atomization) energy by the atom count before scoring. Defaults to True.

Returns:

Maps each target that had at least one value to a (prediction, reference) pair of flat CPU tensors of equal length. Targets with no data are omitted.

Return type:

dict of str to (Tensor, Tensor)