xnn.common.config.schema.DataConfig#

class xnn.common.config.schema.DataConfig(train_path=None, val_path=None, test_path=None, cutoff=4.0, batch_size=16, num_workers=0, val_fraction=0.1, test_fraction=0.0, energy_key='energy', forces_key='forces', stress_key='stress', head=None, replay_path=None, replay_head='pt_head', replay_samples=None, replay_filter='subset', replay_pseudolabel=False)[source]#

Bases: object

Dataset paths and data-loading options.

Variables:
  • train_path (Optional[str]) – Path to the training set (.xyz / .extxyz / .npz / any ASE-readable file). Defaults to None.

  • val_path (Optional[str]) – Path to the validation set. If None, a validation split is carved out of the training set using val_fraction. Defaults to None.

  • test_path (Optional[str]) – Path to a held-out test set, evaluated once after training. If None, a test split is carved out of the training set when test_fraction > 0; otherwise no test evaluation is performed. Defaults to None.

  • cutoff (float) – Neighbor-list cutoff radius; must match model.cutoff and is kept in lockstep by Config.__post_init__(). Defaults to 4.0.

  • batch_size (int) – Mini-batch size; set to 1 to disable batch training. Defaults to 16.

  • num_workers (int) – Number of dataloader worker processes. Defaults to 0.

  • val_fraction (float) – Fraction of the training set held out for validation when val_path is None. Defaults to 0.1.

  • test_fraction (float) – Fraction of the training set held out as a test set when test_path is None. Defaults to 0.0 (no test split).

  • energy_key (str) – Name under which the reference energy is stored in the file (atoms.info), e.g. "REF_energy" for MACE-convention datasets. Defaults to "energy" (also read from the frame’s calculator).

  • forces_key (str) – Name under which the reference forces are stored (atoms.arrays). Defaults to "forces".

  • stress_key (str) – Name under which the reference stress is stored (atoms.info). Defaults to "stress".

  • head (Optional[str]) – For a multi-head model (model.extra["heads"]): the head the training, validation and test structures belong to. None (the default) takes the first head that is not replay_head.

  • replay_path (Optional[str]) – Structures of the pretraining distribution that train the replay head of a multi-head model (multi-head replay fine-tuning). None (default) means no replay set.

  • replay_head (str) – The head the replay structures belong to. Defaults to "pt_head".

  • replay_samples (Optional[int]) – Random subsample size of the replay set (seeded by seed); None (default) keeps every structure that passes the filter.

  • replay_filter (str) – Element filter of the replay set relative to the elements of the training set: "none", "subset" (only training-set elements, the default), "exact" or "superset"; see REPLAY_FILTERS.

  • replay_pseudolabel (bool) – Replace the replay labels by the pretrained model’s own predictions (pseudolabel replay). Defaults to False.

Parameters:
  • train_path (str | None)

  • val_path (str | None)

  • test_path (str | None)

  • cutoff (float)

  • batch_size (int)

  • num_workers (int)

  • val_fraction (float)

  • test_fraction (float)

  • energy_key (str)

  • forces_key (str)

  • stress_key (str)

  • head (str | None)

  • replay_path (str | None)

  • replay_head (str)

  • replay_samples (int | None)

  • replay_filter (str)

  • replay_pseudolabel (bool)