Train a Model

xnn has one configuration schema (the Config dataclass) and three interchangeable frontends to fill it: YAML, argparse, and Hydra.

Write a config file

A YAML training config sets the model, data, and optim sections (any key that is not a core ModelConfig field is folded into model.extra automatically):

# my_train.yaml
model:
  name: mace
  cutoff: 4.0
  species: [18]
  num_channels: 32
  num_interactions: 2
  max_ell: 3
  correlation: 3
  avg_num_neighbors: 20.0

data:
  train_path: data/argon_train.xyz
  val_path: data/argon_val.xyz
  test_path: data/argon_test.xyz   # optional; evaluated once after training
  batch_size: 8

optim:
  lr: 1.0e-3
  epochs: 100
  energy_weight: 1.0
  force_weight: 10.0
  scheduler: plateau

device: auto
seed: 1234
output_dir: runs/argon_mace

The bundled configs/train.yaml additionally shows the Hydra-style defaults: list that composes per-model files from configs/model/ and data settings from configs/data/.

Train from Python

The YAML config can be loaded in Python and passed to the Trainer to initiate the training loop.

from xnn.common.config import from_yaml
from xnn.common.data import AtomicDataset
from xnn.common.train import Trainer

cfg = from_yaml("my_train.yaml")
Trainer(cfg, AtomicDataset(structures, cfg.model.cutoff)).fit()

Train from the command line

The xnn command reads the structure files named in data.train_path / data.val_path with ASE (requires the ase extra) and runs the same trainer:

xnn train --config my_train.yaml

Override any key at the command line

Every frontend supports dotted-key overrides, so a config file can stay generic while runs vary:

xnn train --config my_train.yaml --set optim.epochs=50 model.cutoff=6.0

or from Python:

from xnn.common.config import from_argparse

cfg = from_argparse(["--config", "my_train.yaml", "--set", "model.cutoff=6.0"])

Use Hydra

With the hydra extra installed, an existing Hydra application can hand its DictConfig object straight to xnn’s Trainer:

from xnn.common.config import from_hydra

cfg = from_hydra(hydra_dict_config)