"""Training loop: batch training, device selection, validation, checkpointing.
Keeps the moving parts explicit rather than hiding them in a framework, so the
pipeline is easy to follow and extend.
Multi-GPU / multi-node data parallelism uses native PyTorch DDP and is driven
purely by the environment: when the process was spawned by a distributed
launcher that sets ``RANK`` / ``LOCAL_RANK`` / ``WORLD_SIZE`` (``torchrun
--nproc-per-node N -m xnn train ...``, Slurm + torchrun, or any launcher
that exports the same variables), the trainer initializes the
process group, shards the data loaders with ``DistributedSampler``, wraps the
model in ``DistributedDataParallel``, all-reduces the logged metrics, and
writes checkpoints from rank 0 only. A plain ``python`` / ``xnn`` invocation
runs the unchanged single-process pipeline.
Sharded strategies (FSDP, DeepSpeed ZeRO) are deliberately not used: force and
stress losses back-propagate through gradients taken with
``create_graph=True`` (see ``ForceStressOutput``), a double backward that DDP
supports but sharded wrappers do not.
Fine-tuning a pretrained model is the same loop: ``model.extra["pretrained"]``
starts from any hub model, ``extra["lora"]`` adapts it with low-rank
updates, ``extra["heads"]`` plus a replay set trains a second readout on the
pretraining distribution, and ``atomic_energies: estimated`` re-sets the
per-element references from the training data (see
:mod:`xnn.common.finetune`).
"""
from __future__ import annotations
import contextlib
import copy
import os
import warnings
import torch
from torch.nn.parallel import DistributedDataParallel
from torch.utils.data import ConcatDataset, DataLoader, DistributedSampler, random_split
from .. import distributed
from ..config import Config, TrainerConfig
from ..data import AtomicDataset, collate
from ..models import ForceStressOutput
from ..models.registry import prepare_model
from .callbacks import EarlyStopping, ModelCheckpoint, ProgressBar, monitored_value
from .loggers import build_loggers
from .losses import weighted_loss
[docs]
def resolve_device(name: str) -> torch.device:
"""Resolve a device specification string into a ``torch.device``.
Parameters
----------
name : str
Device name. The special value ``"auto"`` selects ``"cuda"`` when a
CUDA device is available and falls back to ``"cpu"`` otherwise. Any
other value is passed through to ``torch.device`` unchanged
(e.g. ``"cpu"``, ``"cuda"``, ``"cuda:1"``).
Returns
-------
torch.device
The resolved device.
"""
if name == "auto":
return torch.device("cuda" if torch.cuda.is_available() else "cpu")
return torch.device(name)
def _check_charge_labels(base, train_set, optim=None) -> None:
"""Warn when a LatentEwald net-charge constraint will never see a charge label.
With ``constrain_charge`` on, a structure without ``total_charge`` is
pinned to neutrality. A training set in which no structure carries the
key at all is either genuinely neutral, in which case the warning is
harmless, or unlabeled, in which case every charged structure would be
fitted wrong, silently; this makes it visible once, at start. Likewise
the LES charge and dipole terms of ``optim`` are warned about when the
model has no LES wrapper or the set no label for them.
"""
from ..models.les import LatentEwald
les = [m for m in base.modules() if isinstance(m, LatentEwald)]
structures = getattr(train_set, "structures", None)
if optim is not None:
for weight, label in (("les_charge_weight", "charges"), ("les_dipole_weight", "dipole")):
if getattr(optim, weight, 0.0) <= 0:
continue
if not les:
warnings.warn(f"{weight} > 0 but the model has no LES wrapper (long_range): "
"the term is skipped.", stacklevel=3)
elif structures and not any(s.get(label) is not None for s in structures):
warnings.warn(f"{weight} > 0 but no training structure carries a "
f"'{label}' label: the term is skipped.", stacklevel=3)
if not les or not structures:
return
if (any(m.constrain_charge or m.charge_solve for m in les)
and not any(s.get("total_charge", s.get("charge")) is not None for s in structures)):
warnings.warn(
"the model ties its latent charges to each structure's total_charge, "
"but no training structure carries a 'total_charge' (or 'charge') label: all "
f"{len(structures)} are treated as neutral. Label charged structures, or turn "
"constrain_charge / charge_solve off.", stacklevel=3)
if (any(m.fragments for m in les)
and not any(s.get("fragment_charges") is not None for s in structures)):
warnings.warn(
"the charge solve constrains every fragment, but no training structure carries "
"a 'fragment_charges' label: fragment charges come from the ion_charges table "
"alone (monatomic ions), every other fragment is neutral.", stacklevel=3)
def _structures(dataset, what: str) -> list[dict]:
"""The structure dicts behind a dataset (an ``AtomicDataset`` or a list of dicts)."""
if isinstance(dataset, AtomicDataset):
return dataset.structures
if isinstance(dataset, (list, tuple)):
return list(dataset)
raise TypeError(f"{what} must be an AtomicDataset (or a list of structure dicts) for "
f"multi-head training, got {type(dataset).__name__}")
def _set_epoch(loader, epoch: int) -> None:
"""Tell the epoch-aware samplers of a loader which epoch begins.
A ``DistributedSampler`` (``loader.sampler``) or a custom batch sampler
(``loader.batch_sampler``) seeds its shuffle from the epoch and repeats
the same order every epoch unless told; samplers without ``set_epoch``
(the built-in single-process ``RandomSampler`` reshuffles by itself) are
left alone.
Parameters
----------
loader : torch.utils.data.DataLoader
The loader, or any object exposing ``sampler`` / ``batch_sampler``.
epoch : int
The zero-based epoch about to start.
"""
told = []
for sampler in (getattr(loader, "batch_sampler", None), getattr(loader, "sampler", None)):
if hasattr(sampler, "set_epoch") and not any(sampler is t for t in told):
sampler.set_epoch(epoch)
told.append(sampler)
class _EMA:
"""Exponential moving average of a module's parameters, accumulated in float64.
``module`` is a deep copy that holds the averaged weights in the model's own
dtype and is what validation, testing and checkpoints use. The running
average itself lives in float64 shadow tensors: an update adds
``(1 - decay) * (p - avg)`` to the average, and for a parameter that is
large compared with its per-step change (a per-element reference energy
of hundreds of eV moved by Adam steps of 1e-3) that increment is below
float32 resolution. A float32 average would then stay frozen at its first
value while the training copy keeps moving, and the served model would
carry a constant energy offset. The first update copies the parameters,
as ``torch.optim.swa_utils.AveragedModel`` does.
Parameters
----------
module : torch.nn.Module
The training model; its parameters are averaged in registration order.
decay : float
EMA decay in ``(0, 1)``; ``avg = decay * avg + (1 - decay) * p``.
"""
def __init__(self, module: torch.nn.Module, decay: float):
self.module = copy.deepcopy(module)
self.decay = float(decay)
self.n_averaged = 0
self._shadow = [p.detach().to(torch.float64).clone() for p in self.module.parameters()]
@torch.no_grad()
def update_parameters(self, model: torch.nn.Module) -> None:
"""Fold the current parameters of ``model`` into the average."""
d = self.decay if self.n_averaged > 0 else 0.0
for shadow, avg, cur in zip(self._shadow, self.module.parameters(), model.parameters()):
shadow.mul_(d).add_(cur.detach().to(torch.float64), alpha=1.0 - d)
avg.copy_(shadow.to(avg.dtype))
self.n_averaged += 1
#: Properties scored by the error metrics (``trainer.metrics``).
_METRIC_TARGETS = ["energy", "forces", "stress", "dipole", "polarizability", "charges",
"les_dipole", "les_charges"]
class _Sums:
"""Running sums of per-batch ``(logs, counts)`` pairs (see ``Trainer._dataset_metrics``)."""
def __init__(self):
self.sums, self.totals, self.weights, self.n = {}, {}, {}, 0
def add(self, logs: dict, counts: dict) -> "_Sums":
"""Fold in one batch; an RMSE (weight ``None``) is summed as its square."""
self.n += 1
for k, (count, weight) in counts.items():
v = logs[k] ** 2 if weight is None and k.endswith("_rmse") else logs[k]
self.sums[k] = self.sums.get(k, 0.0) + v * count
self.totals[k] = self.totals.get(k, 0.0) + count
self.weights[k] = weight
return self
def _finish(acc: _Sums, is_distributed: bool, device) -> dict:
"""Dataset-level values of accumulated sums, combined over the ranks when distributed."""
sums, totals, weights, n = acc.sums, acc.totals, acc.weights, acc.n
if is_distributed:
weights = distributed.merge_dicts(weights)
keys = sorted(weights)
n, *values = distributed.all_reduce_sum(
[float(n)] + [sums.get(k, 0.0) for k in keys]
+ [totals.get(k, 0.0) for k in keys], device)
sums = dict(zip(keys, values[:len(keys)]))
totals = dict(zip(keys, values[len(keys):]))
if n == 0:
return {}
out = {k: sums[k] / totals[k] for k in sorted(sums) if totals[k] > 0}
loss, heads = 0.0, {}
for k, v in list(out.items()):
if weights[k] is None: # an error metric, not a loss term
if k.endswith("_rmse"):
out[k] = v ** 0.5
continue
loss += weights[k] * v
if "/" in k:
head = k.rsplit("/", 1)[0]
heads[head] = heads.get(head, 0.0) + weights[k] * v
out.update({f"{h}/loss": v for h, v in heads.items()})
out["loss"] = loss
return out
[docs]
class Trainer:
"""Batch training loop with validation, scheduling and checkpointing.
The trainer builds the model from configuration, wraps it in
:class:`ForceStressOutput` (enabling force and/or stress heads only when
the corresponding loss weight is positive), sets up the data loaders,
optimizer and learning-rate scheduler, then runs the full training
pipeline via :meth:`fit`.
When spawned by a distributed launcher (``torchrun -m xnn train ...`` on
one or many nodes; see :meth:`_init_distributed`) the same pipeline runs
data-parallel: the model is wrapped in ``DistributedDataParallel``, each
loader is sharded with a ``DistributedSampler``, metrics are all-reduced
so every rank sees global averages, and only rank 0 logs and writes
checkpoints.
Fine-tuning options of the config are applied while building: the model
section's ``pretrained`` / ``lora`` / ``heads`` entries (see
:func:`~xnn.common.models.registry.prepare_model`), the reference-energy
markers ``atomic_energies: estimated`` / ``average`` (resolved on the
training structures and written back into ``cfg.model`` as numbers),
``optim.freeze`` / ``optim.train_only`` and, for a multi-head model, the
head labels of the datasets and the replay set.
Parameters
----------
cfg : Config
Full run configuration. Fields used include ``device``, ``seed``,
``model`` (model config), ``optim`` (learning rate, weight decay,
epochs, scheduler name and the energy/force/stress loss weights),
``data`` (batch size, number of workers, validation fraction) and
``output_dir`` (where checkpoints are written). ``cfg.model`` is
replaced by the resolved model config (a ``pretrained`` source is
expanded into the architecture it holds), so the saved checkpoints
rebuild the model from their weights alone.
train_set : AtomicDataset
Training dataset of atomic structures.
val_set : AtomicDataset or None, optional
Explicit validation dataset. When ``None`` and
``cfg.data.val_fraction > 0``, a validation split is carved out of
``train_set`` using ``random_split`` seeded by ``cfg.seed``. When
``None`` and the fraction is zero, no validation is performed.
test_set : AtomicDataset or None, optional
Explicit held-out test dataset, evaluated once at the end of
:meth:`fit`. When ``None`` and ``cfg.data.test_fraction > 0``, a test
split is carved out of ``train_set`` (alongside the validation split,
from the same seeded permutation). When ``None`` and the fraction is
zero, no test evaluation is performed.
replay_set : AtomicDataset or None, optional
Structures of the pretraining distribution for the replay head of a
multi-head model (``cfg.data.replay_head``). They are element-filtered
and subsampled as ``cfg.data.replay_filter`` / ``replay_samples`` say,
relabelled by the pretrained model when ``cfg.data.replay_pseudolabel``
is set, split by ``val_fraction`` like the training set, and trained
alongside it. Requires ``model.extra["heads"]``.
train_batch_sampler, val_batch_sampler, test_batch_sampler : iterable or None, optional
Custom batch samplers (each yields lists of dataset indices) for the
training, validation and test loaders, replacing ``cfg.data.batch_size``
and the shuffling, e.g. to keep molecules and periodic cells in
separate batches of their own sizes. A sampler indexes the dataset as
given, so it cannot be combined with a split carved out of
``train_set`` or with a replay set merged into that split. One with a
``set_epoch(epoch)`` method is told each epoch (to reshuffle). In a
distributed run the sampler must yield each rank's own share of every
batch. Defaults to ``None`` (the built-in loaders).
logger : Logger, str, list or False, optional
Where the metrics go: a :class:`~xnn.common.train.loggers.Logger`, a
name (``"console"``, ``"csv"``, ``"tensorboard"``) or a list of
either; ``False`` or ``[]`` logs nothing. ``None`` (the default) uses
``cfg.trainer.loggers``. Only rank 0 logs.
callbacks : list of Callback or None, optional
Hooks into the loop (:mod:`~xnn.common.train.callbacks`). A
:class:`~xnn.common.train.callbacks.ModelCheckpoint` and, as
``cfg.trainer`` asks, an ``EarlyStopping`` and a ``ProgressBar`` are
added unless one of the same type is given.
Attributes
----------
cfg : Config
The configuration passed in, with ``model`` resolved.
distributed : bool
Whether this process is part of a distributed launch (``WORLD_SIZE``
in the environment is greater than one).
rank : int
This process's global rank; 0 in a single-process run.
is_main : bool
Whether this is rank 0, the only rank that logs and saves.
device : torch.device
The resolved training device (``cuda:LOCAL_RANK`` per rank when
distributed on GPUs).
model : ForceStressOutput or DistributedDataParallel
The model wrapped with force/stress output heads, moved to ``device``
(and wrapped in DDP when distributed); :attr:`module` always gives the
bare :class:`ForceStressOutput`.
heads : list of str or None
The head names of a multi-head model, ``None`` otherwise.
ema : _EMA or None
The exponential moving average of the weights when
``cfg.optim.ema_decay > 0``; :attr:`eval_module` is what validation,
testing and the checkpoints use.
train_loader : torch.utils.data.DataLoader
Shuffled loader over the training set (or batched by
``train_batch_sampler``). Batch training is simply ``batch_size > 1``;
set the batch size to 1 to disable it.
val_loader : torch.utils.data.DataLoader or None
Non-shuffled loader over the validation set, or ``None`` if there is
no validation set.
test_loader : torch.utils.data.DataLoader or None
Non-shuffled loader over the test set, or ``None`` if there is no
test set.
opt : torch.optim.Adam or torch.optim.AdamW
The optimizer over the trainable parameters.
sched : torch.optim.lr_scheduler._LRScheduler or ReduceLROnPlateau or None
The learning-rate scheduler, or ``None`` when disabled.
loggers : list of Logger
The loggers :meth:`log_metrics` writes to.
callbacks : list of Callback
The callbacks, in the order they run.
monitor : str
The monitored value (``cfg.trainer.monitor``, with ``val/`` read as
``train/`` when there is no validation set).
mode : str
``"min"`` or ``"max"``.
global_step : int
Optimizer steps taken so far; the ``step`` of every logged value.
should_stop : bool
Set by a callback to end training after the current epoch.
"""
def __init__(self, cfg: Config, train_set: AtomicDataset,
val_set: AtomicDataset | None = None,
test_set: AtomicDataset | None = None,
replay_set: AtomicDataset | None = None,
train_batch_sampler=None, val_batch_sampler=None,
test_batch_sampler=None, logger=None, callbacks=None):
from ..finetune.freeze import freeze_parameters
from ..finetune.heads import find_multihead, label_head
from ..finetune.reference import reference_markers
self.cfg = cfg
# a config unpickled from an older checkpoint has no trainer section
t = getattr(cfg, "trainer", None) or TrainerConfig()
t.validate()
self.tcfg = t
# a batch sampler indexes its dataset as given: refuse the splits and
# the replay merge that would renumber it (before any costly setup)
carves_val = val_set is None and cfg.data.val_fraction > 0
carves_test = test_set is None and cfg.data.test_fraction > 0
if train_batch_sampler is not None and (carves_val or carves_test
or replay_set is not None):
raise ValueError(
"train_batch_sampler indexes train_set as given: pass val_set and "
"test_set explicitly (or set data.val_fraction and data.test_fraction "
"to 0) and no replay_set")
if val_batch_sampler is not None and (val_set is None or replay_set is not None):
raise ValueError("val_batch_sampler needs an explicit val_set and no replay_set")
if test_batch_sampler is not None and test_set is None:
raise ValueError("test_batch_sampler needs an explicit test_set")
self.distributed = distributed.is_distributed()
self.device = self._init_distributed(resolve_device(cfg.device))
self.rank = distributed.rank()
self.is_main = self.rank == 0
torch.manual_seed(cfg.seed)
# reference-energy markers need data: they are resolved below, after
# the model exists, and the config then records the numbers
model_cfg, markers = reference_markers(cfg.model)
base, cfg.model = prepare_model(model_cfg)
# A wrapper such as the D4 dispersion correction can need a larger
# neighbor list than the core model's cutoff that the config carries;
# graphs are built lazily, so widening the datasets' radius here (before
# any graph exists) keeps data and model consistent.
for ds in (train_set, val_set, test_set, replay_set):
model_cutoff = getattr(base, "cutoff", None)
if (ds is not None and model_cutoff is not None
and hasattr(ds, "cutoff") and model_cutoff > ds.cutoff):
ds.cutoff = float(model_cutoff)
if hasattr(ds, "_cache"):
ds._cache.clear()
o = cfg.optim
self._head_weights = {}
multi = find_multihead(base)
self.heads = list(multi.heads) if multi is not None else None
if self.heads:
self._head_weights = self._resolve_head_weights(multi)
force_on = o.force_weight > 0 or any(w[1] > 0 for w in self._head_weights.values())
stress_on = o.stress_weight > 0 or any(w[2] > 0 for w in self._head_weights.values())
self.model = ForceStressOutput(base, compute_forces=force_on,
compute_stress=stress_on).to(self.device)
if self.is_main:
_check_charge_labels(base, train_set, o)
# multi-head: the target sets belong to one head, the replay set to another
replay_train = replay_val = None
if self.heads:
target = self._target_head(multi)
for ds in (train_set, val_set, test_set):
if ds is not None:
label_head(_structures(ds, "the dataset"), multi.index(target))
if replay_set is not None:
replay_train, replay_val = self._prepare_replay(multi, train_set, replay_set)
elif replay_set is not None:
raise ValueError("a replay set needs a multi-head model: set model.heads "
"(for example [pt_head, Default])")
self._resolve_reference_energies(base, multi, markers, train_set)
self._fit_scaling_factors(base, train_set)
if o.freeze or o.train_only:
n_train, n_total = freeze_parameters(self.module, o.freeze, o.train_only)
if self.is_main:
print(f"trainable parameters: {n_train} of {n_total}")
if self.distributed:
self.model = DistributedDataParallel(
self.model,
device_ids=(
[self.device.index] if self.device.type == "cuda" else None),
# a batch without one head leaves that head's parameters unused
find_unused_parameters=bool(self.heads and len(self.heads) > 1))
# Carve val/test splits out of the training set for whichever of the
# two was not given explicitly (a single seeded permutation, so the
# train/val split is unchanged by adding a test fraction of zero).
f_val = cfg.data.val_fraction if val_set is None else 0.0
f_test = cfg.data.test_fraction if test_set is None else 0.0
if f_val > 0 or f_test > 0:
n = len(train_set)
n_val = max(1, int(n * f_val)) if f_val > 0 else 0
n_test = max(1, int(n * f_test)) if f_test > 0 else 0
splits = random_split(
train_set, [n - n_val - n_test, n_val, n_test],
generator=torch.Generator().manual_seed(cfg.seed))
train_set = splits[0]
val_set = splits[1] if n_val else val_set
test_set = splits[2] if n_test else test_set
if replay_train is not None:
train_set = ConcatDataset([train_set, replay_train])
if replay_val is not None:
val_set = replay_val if val_set is None else ConcatDataset([val_set, replay_val])
# batch training is just batch_size > 1; set to 1 to disable.
# Distributed runs shard every loader across ranks. Training uses
# DistributedSampler, which owns the shuffling (the two DataLoader
# options are exclusive) and pads the shards to equal length so every
# rank takes the same number of synchronised steps. Evaluation has no
# gradient sync, so its shards are disjoint and unpadded: no structure
# is counted twice in the validation and test metrics.
def _loader(ds, shuffle, batch_sampler=None):
if batch_sampler is not None:
return DataLoader(ds, batch_sampler=batch_sampler, collate_fn=collate,
num_workers=cfg.data.num_workers)
if not self.distributed:
sampler = None
elif shuffle:
sampler = DistributedSampler(ds, shuffle=True, seed=cfg.seed)
else:
sampler = distributed.shard_indices(len(ds))
return DataLoader(ds, batch_size=cfg.data.batch_size,
shuffle=shuffle and sampler is None,
sampler=sampler, collate_fn=collate,
num_workers=cfg.data.num_workers)
self.train_loader = _loader(train_set, True, train_batch_sampler)
self.val_loader = (_loader(val_set, False, val_batch_sampler)
if val_set is not None else None)
self.test_loader = (_loader(test_set, False, test_batch_sampler)
if test_set is not None else None)
self._params = [p for p in self.model.parameters() if p.requires_grad]
if not self._params:
raise ValueError("no trainable parameters: every parameter is frozen")
name = o.optimizer.lower()
optimizer = {"adam": torch.optim.Adam, "amsgrad": torch.optim.Adam,
"adamw": torch.optim.AdamW}[name]
self.opt = optimizer(self._params, lr=o.lr, weight_decay=o.weight_decay,
amsgrad=(name == "amsgrad"))
self.sched = self._make_scheduler(o.scheduler)
self.ema = None
if o.ema_decay > 0:
if not 0.0 < o.ema_decay < 1.0:
raise ValueError(f"optim.ema_decay must lie in (0, 1), got {o.ema_decay}")
self.ema = _EMA(self.module, o.ema_decay)
os.makedirs(cfg.output_dir, exist_ok=True)
self.global_step = 0
self.should_stop = False
self.loggers = build_loggers(t.loggers if logger is None else logger)
for lg in self.loggers:
lg.setup(cfg.output_dir)
self.monitor, self.mode = t.monitor, t.mode
if self.val_loader is None and self.monitor.startswith("val/"):
self.monitor = "train/" + self.monitor[len("val/"):]
self.callbacks = list(callbacks or [])
defaults = [ModelCheckpoint(every_n_epochs=t.save_every_n_epochs)]
if t.early_stopping_patience > 0 or t.early_stopping_threshold is not None:
defaults.append(EarlyStopping(patience=t.early_stopping_patience,
min_delta=t.early_stopping_min_delta,
stopping_threshold=t.early_stopping_threshold))
if t.progress_bar:
defaults.append(ProgressBar())
for cb in defaults:
if not any(isinstance(c, type(cb)) for c in self.callbacks):
self.callbacks.append(cb)
def _init_distributed(self, dev: torch.device) -> torch.device:
"""Join the process group of a distributed launcher, if there is one.
A launcher such as ``torchrun`` exports
``RANK`` / ``LOCAL_RANK`` / ``WORLD_SIZE`` into every process it
spawns; ``self.distributed`` reflects whether that happened. In a
distributed run each rank is pinned to one CUDA device selected by
``LOCAL_RANK`` (the configured device only chooses cpu vs cuda), and
the process group is initialized with the matching backend -- NCCL on
GPUs, Gloo on CPUs. The work is :func:`xnn.common.distributed.init_process_group`,
shared with the benchmark runner.
Parameters
----------
dev : torch.device
The device resolved from the configuration.
Returns
-------
torch.device
The device this rank should train on: ``dev`` unchanged in a
single-process run, ``cuda:LOCAL_RANK`` (or ``dev`` on CPU) in a
distributed one.
"""
dev, self._owns_pg = distributed.init_process_group(dev)
return dev
@property
def module(self) -> ForceStressOutput:
"""The bare model, unwrapped from ``DistributedDataParallel`` if any."""
if isinstance(self.model, DistributedDataParallel):
return self.model.module
return self.model
@property
def eval_module(self) -> ForceStressOutput:
"""The weights used for validation, testing and checkpoints.
The exponential moving average when ``cfg.optim.ema_decay > 0``
(its buffers synced from the training model), otherwise
:attr:`module`.
"""
if self.ema is None:
return self.module
with torch.no_grad():
for avg, cur in zip(self.ema.module.buffers(), self.module.buffers()):
if avg.shape == cur.shape:
avg.copy_(cur)
return self.ema.module
def _target_head(self, multi) -> str:
"""The head the training structures belong to (``cfg.data.head`` or the first non-replay one)."""
if self.cfg.data.head is not None:
return self.cfg.data.head
for name in multi.heads:
if name != self.cfg.data.replay_head:
return name
return multi.heads[0]
def _resolve_head_weights(self, multi) -> dict[int, tuple[float, float, float]]:
"""``cfg.optim.head_weights`` as ``{head_index: (energy, force, stress)}``."""
o = self.cfg.optim
out = {}
for name, spec in (o.head_weights or {}).items():
spec = dict(spec or {})
unknown = sorted(set(spec) - {"energy_weight", "force_weight", "stress_weight"})
if unknown:
raise ValueError(f"optim.head_weights[{name!r}]: unknown key(s) {unknown}")
out[multi.index(str(name))] = (
float(spec.get("energy_weight", o.energy_weight)),
float(spec.get("force_weight", o.force_weight)),
float(spec.get("stress_weight", o.stress_weight)))
return out
def _prepare_replay(self, multi, train_set, replay_set):
"""Filter, subsample, (pseudo)label and split the replay set.
Returns
-------
tuple
``(replay_train, replay_val)`` datasets labelled with the replay
head (``replay_val`` is ``None`` without a validation fraction).
"""
from ..finetune.heads import label_head
from ..finetune.reference import species_of
from ..finetune.replay import pseudolabel, select_replay
d = self.cfg.data
structures = _structures(replay_set, "replay_set")
species = species_of(_structures(train_set, "train_set")) if d.replay_filter != "none" else None
selected = select_replay(structures, species, n=d.replay_samples, mode=d.replay_filter,
seed=self.cfg.seed)
if not selected:
raise ValueError("the replay set is empty after the element filter "
f"({d.replay_filter}); use replay_filter: none or another set")
if d.replay_pseudolabel:
selected = pseudolabel(multi, selected, batch_size=d.batch_size, device=self.device,
head=d.replay_head)
label_head(selected, multi.index(d.replay_head))
if self.is_main:
print(f"replay head {d.replay_head!r}: {len(selected)} structures"
f"{' (pseudolabelled)' if d.replay_pseudolabel else ''}")
cutoff = float(getattr(replay_set, "cutoff", self.cfg.data.cutoff))
ds = AtomicDataset(selected, cutoff)
if d.val_fraction > 0 and len(ds) > 1:
n_val = max(1, int(len(ds) * d.val_fraction))
train, val = random_split(ds, [len(ds) - n_val, n_val],
generator=torch.Generator().manual_seed(self.cfg.seed + 1))
return train, val
return ds, None
def _fit_scaling_factors(self, base, train_set) -> None:
"""Fit the activation scaling factors of a model that has unfitted ones (GemNet).
25 batches of 32 training structures, as in the GemNet reference code;
no labels are used. ``model.extra["fit_scaling"] = False`` skips it.
"""
targets = [m for m in base.modules()
if callable(getattr(m, "fit_scaling_factors", None))
and not bool(getattr(m, "scaling_fitted", True))]
if not targets or not bool((self.cfg.model.extra or {}).get("fit_scaling", True)):
return
loader = DataLoader(train_set, batch_size=32, shuffle=True, collate_fn=collate,
generator=torch.Generator().manual_seed(self.cfg.seed))
batches = [graph.to(self.device) for _, graph in zip(range(25), loader)]
for module in targets:
fitted = module.fit_scaling_factors(batches)
if self.is_main:
print(f"fitted {len(fitted)} scaling factors on "
f"{sum(g.num_graphs for g in batches)} structures")
def _resolve_reference_energies(self, base, multi, markers, train_set) -> None:
"""Estimate the reference energies the config marked, set them, record them."""
if not markers:
return
from ..finetune.heads import head_options
from ..finetune.reference import (average_atomic_energies, estimate_atomic_energies,
get_atomic_energies, set_atomic_energies)
structures = _structures(train_set, "train_set")
extra = dict(self.cfg.model.extra or {})
# the recorded values cover every species of the model (a per-species
# list in the config must be complete), the estimate only those present
def complete(single, e0):
full = get_atomic_energies(single)
full.update(e0)
return full
heads_spec = extra.get("heads")
if multi is not None and isinstance(heads_spec, (list, tuple)):
heads_spec = {name: {} for name in heads_spec}
for head, marker in markers:
if multi is not None:
head = head if head is not None else self._target_head(multi)
target_structures = [s for s in structures
if int(s.get("head", 0)) == multi.index(head)]
scope = multi.using(head)
else:
target_structures = structures
scope = contextlib.nullcontext(base)
with scope as single:
if marker == "estimated":
e0 = estimate_atomic_energies(single, target_structures,
batch_size=self.cfg.data.batch_size,
device=self.device)
else:
e0 = average_atomic_energies(target_structures)
set_atomic_energies(single, e0)
recorded = complete(single, e0)
if self.is_main:
where = f"head {head!r}" if multi is not None else "the model"
print(f"{marker} reference energies for {where}: "
+ ", ".join(f"{z}: {v:.4f}" for z, v in e0.items()))
if multi is not None:
heads_spec[head] = {**head_options(heads_spec, head), "atomic_energies": recorded}
else:
extra["atomic_energies"] = recorded
if multi is not None:
extra["heads"] = heads_spec
self.cfg.model.extra = extra
def _make_scheduler(self, name: str):
"""Construct the learning-rate scheduler named in the config.
Parameters
----------
name : str
Scheduler identifier. ``"cosine"`` builds a
``CosineAnnealingLR`` over ``cfg.optim.epochs``; ``"plateau"``
builds a ``ReduceLROnPlateau`` with ``patience=10``. Any other
value disables scheduling.
Returns
-------
torch.optim.lr_scheduler.CosineAnnealingLR or torch.optim.lr_scheduler.ReduceLROnPlateau or None
The scheduler instance, or ``None`` when scheduling is disabled.
"""
if name == "cosine":
return torch.optim.lr_scheduler.CosineAnnealingLR(
self.opt, T_max=self.cfg.optim.epochs)
if name == "plateau":
return torch.optim.lr_scheduler.ReduceLROnPlateau(self.opt, mode=self.tcfg.mode,
patience=10)
return None
def _step(self, data, train: bool):
"""Run one forward/loss step on a batch, optionally back-propagating.
Moves the batch to the training device, runs the model forward,
computes the weighted energy/force/stress loss, and (when ``train``)
performs a single optimizer step (with gradient clipping and the EMA
update when configured).
Parameters
----------
data : AtomicGraph
A (possibly batched) atomic graph produced by the collate
function.
train : bool
When ``True``, zero the gradients, back-propagate the loss and
step the optimizer. When ``False`` (validation), only the forward
pass and loss computation are performed on :attr:`eval_module`;
note that gradients are still enabled because force predictions
require them.
Returns
-------
tuple of (dict of str to float, dict)
The scalar loss logs for this batch, as returned by
:func:`weighted_loss` (e.g. ``"loss"`` plus any of
``"energy_mse"``, ``"force_mse"``, ``"stress_mse"``,
``"dipole_mse"``, ``"polarizability_mse"``, ``"charge_mse"``, per head
in a multi-head run, plus the error metrics such as
``"forces_mae"``), and the ``counts`` of its terms that
:meth:`_dataset_metrics` accumulates.
"""
data = data.to(self.device)
o = self.cfg.optim
# Evaluation steps run the bare module (the EMA weights when kept):
# they are not followed by a backward pass, so the DDP wrapper's
# gradient sync must not be armed.
model = self.model if train else self.eval_module
pred = model(data)
counts = {}
loss, logs = weighted_loss(
pred, data, o.energy_weight, o.force_weight, o.stress_weight,
dipole_weight=o.dipole_weight,
polarizability_weight=o.polarizability_weight,
charge_weight=o.charge_weight,
les_dipole_weight=o.les_dipole_weight,
les_charge_weight=o.les_charge_weight,
huber_delta=o.huber_delta,
huber_delta_energy=o.huber_delta_energy,
huber_delta_forces=o.huber_delta_forces,
huber_delta_stress=o.huber_delta_stress,
head_weights=self._head_weights or None, head_names=self.heads,
counts=counts)
if self.tcfg.metrics:
self._error_metrics(pred, data, logs, counts)
if train:
if self.distributed:
# an output without labels in this batch (a dipole head on a
# batch without dipoles) still joins the graph with weight 0,
# so DDP sees every parameter used; values and gradients are
# unchanged
loss = loss + sum(0.0 * v.sum() for k, v in pred.items()
if k not in ("forces", "stress", "virials")
and torch.is_tensor(v) and v.requires_grad)
self.opt.zero_grad()
loss.backward()
if o.clip_grad > 0:
torch.nn.utils.clip_grad_norm_(self._params, o.clip_grad)
self.opt.step()
self.global_step += 1
if self.ema is not None:
self.ema.update_parameters(self.module)
return logs, counts
def _error_metrics(self, pred, data, logs: dict, counts: dict) -> None:
"""Add the MAE / RMSE of every labelled property of a batch to ``logs`` and ``counts``.
The pairs are those the benchmark scores
(:func:`~xnn.common.benchmark.metrics.prediction_pairs`: energy per
atom, every component of forces, stress, dipole and polarizability,
charges per atom, zero-weight labels left out), unweighted. Their
counts carry the weight ``None``, so they stay out of the loss.
"""
from ..benchmark.metrics import prediction_pairs
for target, (p, r) in prediction_pairs(pred, data, _METRIC_TARGETS).items():
n = p.numel()
if n == 0:
continue
d = (p - r).detach().double()
if "mae" in self.tcfg.metrics:
logs[f"{target}_mae"] = float(d.abs().mean())
counts[f"{target}_mae"] = (float(n), None)
if "rmse" in self.tcfg.metrics:
logs[f"{target}_rmse"] = float(d.pow(2).mean().sqrt())
counts[f"{target}_rmse"] = (float(n), None)
def _callback(self, hook: str, *args) -> None:
"""Run ``hook`` of every callback."""
for cb in self.callbacks:
getattr(cb, hook)(self, *args)
def _train_epoch(self) -> dict:
"""One training pass over the loader; returns its dataset-level metrics.
With ``cfg.trainer.log_every_n_steps > 0`` the training values are
logged every n steps as ``train_step/<key>``: the step's batch or,
with ``step_average: running``, the epoch so far, combined over the
ranks like the epoch values.
"""
t = self.tcfg
epoch = _Sums()
for i, data in enumerate(self.train_loader):
logs, counts = self._step(data, True)
epoch.add(logs, counts)
self._callback("on_train_batch_end", logs, i)
if t.log_every_n_steps > 0 and self.global_step % t.log_every_n_steps == 0:
acc = epoch if t.step_average == "running" else _Sums().add(logs, counts)
values = _finish(acc, self.distributed, self.device)
self.log_metrics({f"train_step/{k}": v for k, v in values.items()})
return _finish(epoch, self.distributed, self.device)
[docs]
def log_metrics(self, metrics: dict, step: int | None = None) -> None:
"""Send ``metrics`` to every logger (rank 0 only).
Parameters
----------
metrics : dict of str to float
Metric names mapped to values.
step : int or None, optional
The step to record them at; defaults to :attr:`global_step`.
"""
if not self.is_main:
return
step = self.global_step if step is None else step
for lg in self.loggers:
lg.log_metrics(metrics, step)
[docs]
def fit(self):
"""Run the training loop over ``cfg.optim.epochs`` epochs (or until a callback stops it).
Each epoch trains over the whole training loader, validates every
``cfg.trainer.val_every_n_epochs`` epochs (and after the last), steps
the scheduler (a ``ReduceLROnPlateau`` with the monitored value, on
the epochs that have it) and runs the callbacks: the default
:class:`~xnn.common.train.callbacks.ModelCheckpoint` writes
``best.pt`` whenever the monitored value improves and ``last.pt``
after the last epoch. The epoch metrics go to the loggers every
``cfg.trainer.log_every_n_epochs`` epochs (``epoch``, ``lr``,
``train/<key>``, ``val/<key>``), and the test metrics, evaluated once
with the final weights, as ``test/<key>``. Every value is a
dataset-level mean (see :meth:`_dataset_metrics`). The training
loader's sampler is told the epoch before each pass (``set_epoch``)
whenever it has that method, in every launch mode. To test the best
checkpoint instead, load ``best.pt`` and call :meth:`evaluate`.
Returns
-------
dict
Final metrics: ``{"train": ..., "val": ..., "test": ...}``, each a
dict of dataset-level values (empty when that split does not exist).
The training metrics are accumulated during the epoch, while the
weights change.
"""
if self.is_main and self.loggers:
from ..models.hub.checkpoint import _dataclass_dict
params = _dataclass_dict(self.cfg) # every field, defaults included
for lg in self.loggers:
lg.log_hyperparams(params)
status = "failed"
try:
out = self._fit()
status = "success"
return out
finally:
if self.is_main:
for lg in self.loggers:
lg.finalize(status)
def _fit(self):
"""The epoch loop of :meth:`fit`."""
t = self.tcfg
epochs = self.cfg.optim.epochs
self.should_stop = False
self._callback("on_fit_start")
tr, va = {}, {}
for epoch in range(epochs):
last = epoch == epochs - 1
_set_epoch(self.train_loader, epoch)
lr = self.opt.param_groups[0]["lr"]
self._callback("on_train_epoch_start", epoch)
self.model.train()
tr = self._train_epoch()
va = {}
if self.val_loader is not None and ((epoch + 1) % t.val_every_n_epochs == 0 or last):
self.model.eval()
# forces need grad even at eval -> no torch.no_grad()
va = self._dataset_metrics(self._step(d, False) for d in self.val_loader)
metrics = {"epoch": epoch, "lr": lr, **{f"train/{k}": v for k, v in tr.items()},
**{f"val/{k}": v for k, v in va.items()}}
if isinstance(self.sched, torch.optim.lr_scheduler.ReduceLROnPlateau):
value = monitored_value(metrics, self.monitor)
if value is not None:
self.sched.step(value)
elif self.sched is not None:
self.sched.step()
self._callback("on_epoch_end", epoch, metrics)
if (epoch + 1) % t.log_every_n_epochs == 0 or last or self.should_stop:
self.log_metrics(metrics)
if self.should_stop:
break
self._callback("on_train_end")
te = {}
if self.test_loader is not None:
te = self.evaluate()
self.log_metrics({f"test/{k}": v for k, v in te.items()})
distributed.destroy_process_group(self._owns_pg)
results = {"train": tr, "val": va, "test": te}
self._callback("on_fit_end", results)
return results
[docs]
def evaluate(self, loader=None):
"""Evaluate the current model over a data loader without training.
In a distributed run each rank scores its shard and the metrics are
combined across ranks. After :meth:`fit` has ended the process group
(the trainer owned it), every rank scores the whole set of the loader
instead, so the metrics are still those of the whole set.
Parameters
----------
loader : torch.utils.data.DataLoader or None, optional
Loader to evaluate over. Defaults to ``self.test_loader``.
Returns
-------
dict of str to float
Dataset-level losses over the loader (see
:meth:`_dataset_metrics`), or an empty dict when there is no
loader.
"""
loader = self.test_loader if loader is None else loader
if loader is None:
return {}
if self.distributed and not (torch.distributed.is_available()
and torch.distributed.is_initialized()):
# fit() has ended the process group, so the ranks can no longer
# combine their shards: every rank scores the whole set instead
loader = DataLoader(loader.dataset, batch_size=self.cfg.data.batch_size,
collate_fn=collate, num_workers=self.cfg.data.num_workers)
self.model.eval()
# forces need grad even at eval -> no torch.no_grad()
return self._dataset_metrics(self._step(d, False) for d in loader)
def _dataset_metrics(self, steps):
"""Accumulate per-batch losses into means over the whole dataset.
Each loss term of a batch is a weighted mean over its structures or
atoms, so the batches are combined by those weights, not counted
equally: a term's dataset value is ``sum(mean * count) / sum(count)``
with the counts of :func:`weighted_loss`, which equals the term over
one batch holding every structure. A smaller last batch, or a term
that only some batches carry (stress of the periodic structures,
a head's own structures), is then weighted by what it holds. The
``"loss"`` is rebuilt from those means and the term weights, and so
is the ``"<head>/loss"`` of each head. The error metrics are pooled
the same way over every scored value (an RMSE through its square).
In a distributed run the per-rank sums are all-reduced over the union
of every rank's keys, so every rank returns the same global values
and the best-checkpoint decision and the plateau scheduler stay in
lockstep across ranks.
Parameters
----------
steps : iterable of tuple
The ``(logs, counts)`` pairs returned by :meth:`_step`.
Returns
-------
dict of str to float
Each term key mapped to its dataset-level mean, plus ``"loss"``
(and the per-head losses). An empty iterable on every rank yields
an empty dictionary.
"""
acc = _Sums()
for logs, counts in steps:
acc.add(logs, counts)
return _finish(acc, self.distributed, getattr(self, "device", None))
[docs]
def save(self, path: str):
"""Serialize the model state dict and configuration to disk.
Only rank 0 writes (in a single-process run that is the only rank);
the state dict is taken from the bare module (the EMA weights when
kept), so checkpoint keys are identical with and without DDP. A
barrier keeps the other ranks from racing ahead of the write.
Parameters
----------
path : str
Destination file path. The saved checkpoint is a dictionary with
keys ``"model"`` (the model ``state_dict``) and ``"cfg"`` (the
:class:`Config` used for the run, with the resolved model
section).
Returns
-------
None
"""
if self.is_main:
torch.save({"model": self.eval_module.state_dict(), "cfg": self.cfg}, path)
distributed.barrier()
[docs]
def save_pretrained(self, save_directory: str, **card_fields):
"""Write the current model as a portable model directory.
The layout :func:`~xnn.common.models.hub.from_pretrained` loads and
the model hub caches (``card.json``, ``config.yaml``, ``model.pt``),
with no pickled objects, so the directory can be shared, uploaded to
Zenodo, or copied to another machine. Only rank 0 writes. A
multi-head or LoRA model is written as such (its config rebuilds it);
``from_pretrained`` then serves one head with the LoRA updates folded
in.
Parameters
----------
save_directory : str
Output directory.
**card_fields
Model card fields (``name``, ``description``, ``license``,
``citation``, ...); see :class:`~xnn.common.models.hub.ModelCard`.
Returns
-------
pathlib.Path or None
The directory on rank 0, ``None`` on the other ranks.
"""
from ..models.hub import save_pretrained
out = None
if self.is_main:
out = save_pretrained(self.eval_module, save_directory, config=self.cfg, **card_fields)
distributed.barrier()
return out
