Ctrl+K

nvalchemi.training.EnergyMSELoss#

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

Mean-squared-error loss on per-graph total energy.

Energies enter this loss as one total-energy value per graph, with canonical shape (B, 1). With per_atom=False the loss is the graph-balanced MSE of total-energy residuals, so every graph has equal weight regardless of size:

\[L = \frac{1}{B} \sum_{i=1}^{B} \left(\hat{E}_i - E_i\right)^2.\]

With per_atom=True the prediction and target are first divided by each graph’s atom count \(N_i\), so the residual is measured in energy-per-atom units, and the reduction is weighted by \(N_i\) 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)^2}{\sum_{i=1}^{B} N_i} = \frac{\sum_{i=1}^{B} (\hat{E}_i - E_i)^2 / N_i}{\sum_{i=1}^{B} N_i}.\]

A hat denotes the prediction and \(N_i\) the atom count of graph \(i\) (see the module docstring for the shared notation). Counts \(N_i\) may be supplied directly as (B,) or recovered from a padded node mask of shape (B, V_max).

Tensor Contract#

pred, targetEnergy

Per-graph energy tensors of shape (B, 1). Shape validation requires exact equality; (B, 1) and (B,) are rejected even though they are broadcast-compatible.

num_nodes_per_graphInteger[torch.Tensor, “B”] | Bool[torch.Tensor, “B V_max”], optional

Required only when per_atom=True. May be explicit per-graph counts or a padded node-validity mask.

param target_key:

Target container key for the target tensor.

type target_key:

str, default “energy”

param prediction_key:

Prediction container key for the model output.

type prediction_key:

str, default “predicted_energy”

param per_atom:

Measure residuals in energy-per-atom units and reduce them with atom-count weights: larger graphs contribute in proportion to their atom counts.

type per_atom:

bool, default False

param ignore_nonfinite:

When True, target entries that are NaN or infinite are excluded from both loss value and gradient using torch.isfinite(). Intended for inputs where some samples lack an energy label. Implemented with branch-free tensor ops for torch.compile compatibility. When per_atom=True, atom-count weights for invalid targets are also excluded from the denominator. When every target entry is non-finite the loss is 0.0.

type ignore_nonfinite:

bool, default False

param dtype_policy:

How to handle prediction/target dtype mismatches before validation. "strict" raises; the other policies cast one tensor to match the other.

type dtype_policy:

{“strict”, “prediction_to_target”, “target_to_prediction”}, default “strict”

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

Return squared 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]

Parameters:

  • target_key (str)

  • prediction_key (str)

  • per_atom (bool)

  • ignore_nonfinite (bool)

  • dtype_policy (DTypePolicy)

On this page
Show Source