Ctrl+K

nvalchemi.training.EnergyMAELoss#

class nvalchemi.training.EnergyMAELoss(*, target_key='energy', prediction_key='predicted_energy', per_atom=True, ignore_nonfinite=True, dtype_policy='strict')[source]#

Mean-absolute-error loss for per-graph energy targets.

This loss operates on per-graph total energies with identical prediction and target shapes, commonly (B, 1) or (B,). With per_atom=True (default), prediction and target energies are first divided by each graph’s atom count \(N_i\), then absolute residuals are reduced with atom-count weights so that larger graphs contribute in proportion to their size:

\[L = \frac{\sum_{i=1}^{B} N_i \left|\dfrac{\hat{E}_i - E_i}{N_i}\right|}{\sum_{i=1}^{B} N_i} = \frac{\sum_{i=1}^{B} |\hat{E}_i - E_i|}{\sum_{i=1}^{B} N_i}.\]

With per_atom=False the loss is the graph-balanced mean absolute error of total-energy residuals, \(L = \tfrac{1}{B} \sum_{i=1}^{B} |\hat{E}_i - E_i|\). A hat denotes the prediction and \(N_i\) the atom count of graph \(i\) (see the module docstring for the shared notation).

Parameters:

  • target_key (str, default "energy") – Target container key for the target tensor.

  • prediction_key (str, default "predicted_energy") – Prediction container key for the model output.

  • per_atom (bool, default True) – Divide prediction and target by num_nodes_per_graph before computing absolute residuals.

  • ignore_nonfinite (bool, default True) – When True, target entries that are NaN or infinite are excluded from both loss value and gradient using torch.isfinite().

  • dtype_policy ({"strict", "prediction_to_target", "target_to_prediction"}, default "strict") – How to handle prediction/target dtype mismatches before validation. "strict" raises; the other policies cast one tensor to match the other.

compute_residual(pred, target, valid)[source]#

Return absolute residuals, zeroing invalid entries.

Parameters:
Return type:

Tensor

extra_repr()[source]#

Human-readable hyperparameter summary for nn.Module’s repr.

Return type:

str

mask(pred, target, ctx, **kwargs)[source]#

Exclude non-finite target entries when ignore_nonfinite=True.

Parameters:
Return type:

Tensor

normalize(pred, target, **kwargs)[source]#

Divide by atom counts when per_atom=True.

Parameters:
Return type:

tuple[Tensor, Tensor, ReductionContext]

On this page
Show Source