xnn.common.train.trainer.Trainer#

class xnn.common.train.trainer.Trainer(cfg, train_set, val_set=None, test_set=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.

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

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

Variables:

  • cfg (Config) – The configuration passed in.

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

  • train_loader (torch.utils.data.DataLoader) – Shuffled loader over the training set. 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) – The Adam optimizer.

  • sched (torch.optim.lr_scheduler._LRScheduler or ReduceLROnPlateau or None) – The learning-rate scheduler, or None when disabled.

property module: ForceStressOutput#

The bare model, unwrapped from DistributedDataParallel if any.

fit()[source]#

Run the full training loop over cfg.optim.epochs epochs.

For each epoch the model is trained over the whole training loader and, if a validation loader exists, evaluated over it. The best validation loss seen so far is tracked and its model saved to best.pt in the output directory. The scheduler is stepped every epoch: a ReduceLROnPlateau scheduler is stepped with the current loss, while any other scheduler is stepped without arguments. A per-epoch summary is printed. After the final epoch the model is saved to last.pt and, if a test loader exists, evaluated once on the test set (with the final-epoch weights; to test the best checkpoint instead, load best.pt into self.model and call evaluate()).

Returns:

Final metrics: {"train": ..., "val": ..., "test": ...}, each a dict of averaged losses (empty when that split does not exist).

Return type:

dict

evaluate(loader=None)[source]#

Evaluate the current model over a data loader without training.

Parameters:

loader (torch.utils.data.DataLoader or None, optional) – Loader to evaluate over. Defaults to self.test_loader.

Returns:

Averaged losses over the loader (as in _step()), 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, 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).

Return type:

None