xnn.common.train.trainer.Trainer#

class xnn.common.train.trainer.Trainer(cfg, train_set, val_set=None, test_set=None, replay_set=None, train_batch_sampler=None, val_batch_sampler=None, test_batch_sampler=None, logger=None, callbacks=None)[source]#

Bases: object

Batch training loop with validation, scheduling and checkpointing.

The trainer builds the model from configuration, wraps it in 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 fit().

When spawned by a distributed launcher (torchrun -m xnn train ... on one or many nodes; see _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 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 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 (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).

  • val_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).

  • 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 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 (callbacks). A ModelCheckpoint and, as cfg.trainer asks, an EarlyStopping and a ProgressBar are added unless one of the same type is given.

Variables:

  • 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); module always gives the bare 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; 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 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.

property module: ForceStressOutput#

The bare model, unwrapped from DistributedDataParallel if any.

property eval_module: 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 module.

log_metrics(metrics, step=None)[source]#

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 global_step.

Return type:

None

fit()[source]#

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 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 _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 evaluate().

Returns:

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.

Return type:

dict

evaluate(loader=None)[source]#

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 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:

Dataset-level losses over the loader (see _dataset_metrics()), or an empty dict when there is no loader.

Return type:

dict of str to float

save(path)[source]#

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 Config used for the run, with the resolved model section).

Return type:

None

save_pretrained(save_directory, **card_fields)[source]#

Write the current model as a portable model directory.

The layout 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 ModelCard.

Returns:

The directory on rank 0, None on the other ranks.

Return type:

pathlib.Path or None