xnn.common.config.schema.OptimConfig#

class xnn.common.config.schema.OptimConfig(lr=0.001, weight_decay=0.0, epochs=100, energy_weight=1.0, force_weight=10.0, stress_weight=0.0, dipole_weight=0.0, polarizability_weight=0.0, charge_weight=0.0, les_dipole_weight=0.0, les_charge_weight=0.0, scheduler='plateau', huber_delta=0.0, huber_delta_energy=None, huber_delta_forces=None, huber_delta_stress=None, optimizer='adam', clip_grad=0.0, ema_decay=0.0, head_weights=None, freeze=<factory>, train_only=<factory>)[source]#

Bases: object

Optimizer, schedule, and loss-weighting settings.

Variables:

  • lr (float) – Learning rate. Defaults to 1e-3.

  • weight_decay (float) – L2 weight-decay coefficient. Defaults to 0.0.

  • epochs (int) – Number of training epochs. MACE max_num_epochs. Defaults to 100.

  • energy_weight (float) – Weight of the energy term in the loss. MACE energy_weight. Defaults to 1.0.

  • force_weight (float) – Weight of the force term in the loss. Defaults to 10.0.

  • stress_weight (float) – Weight of the stress term; values > 0 enable stress training for periodic systems. Defaults to 0.0.

  • dipole_weight (float) – Weight of the dipole-moment term, for models with a dipole output (PaiNN, PhysNet, AIMNet2) on data with dipole labels. Defaults to 0.0.

  • polarizability_weight (float) – Weight of the polarizability-tensor term (PaiNN). Defaults to 0.0.

  • charge_weight (float) – Weight of the per-atom partial-charge term, for models with a charges output (the 3G / 4G HDNNP charge networks, PhysNet, AIMNet2) on data with charges labels. Defaults to 0.0.

  • les_dipole_weight (float) – Weight of the LES dipole term: the les_dipole of a model wrapped in LatentEwald against the dipole labels. Defaults to 0.0.

  • les_charge_weight (float) – Weight of the LES charge term: les_charges against the charges labels. Defaults to 0.0.

  • scheduler (str) – Learning-rate scheduler (none / cosine / plateau). MACE uses ReduceLROnPlateau. Defaults to "plateau".

  • huber_delta (float) – Crossover from quadratic to linear in the loss, which caps the pull of a few large residuals. 0.0 (the default) is plain squared error. Note xnn scales the Huber function to agree with the squared error below delta, so the loss weights keep their meaning when it is switched on; MACE uses the textbook half-square form.

  • huber_delta_energy (float, optional) – Per-term override of huber_delta for the energy term. None (default) falls back to it. Energies (eV per atom) and forces (eV/A) differ in scale by an order of magnitude, so one delta rarely suits both.

  • huber_delta_forces (float, optional) – Per-term override of huber_delta for the force term.

  • huber_delta_stress (float, optional) – Per-term override of huber_delta for the stress term.

  • optimizer (str) – "adam" (default; weight_decay is an L2 penalty), "amsgrad" (Adam with the AMSGrad maximum of the second moment, whose steps shrink as training settles; the DimeNet papers’ choice) or "adamw" (decoupled weight decay, the fine-tuning protocols’ choice).

  • clip_grad (float) – Maximum gradient norm per step (torch.nn.utils.clip_grad_norm_); 0.0 (default) disables clipping.

  • ema_decay (float) – Decay of an exponential moving average of the weights, updated after every optimizer step and used for validation, testing and the saved checkpoints (0.99-0.9999 in the fine-tuning protocols). 0.0 (default) trains and saves the raw weights.

  • head_weights (dict, optional) – Per-head loss weights of a multi-head model, e.g. {"pt_head": {"energy_weight": 1.0, "force_weight": 10.0}}; a head not listed uses the global weights. Each head’s terms are averaged over its own structures and the heads are summed.

  • freeze (list of str) – fnmatch patterns of parameter names to freeze (layer-freezing fine-tuning), e.g. ["model.node_embedding*", "model.interactions.0.*"].

  • train_only (list of str) – Patterns of the only parameters to train; everything else is frozen (["*readouts*", "*atom_ref*"] trains the readout alone).

Parameters:
On this page
Show Source