Scaler model

metatrain.scaler.train_or_load_scaler(scaler: Scaler, train_datasets: List[Dataset | Subset], additive_models: List[Module], batch_size: int, is_distributed: bool, fixed_weights: Dict[str, float | Dict[int, float]] | str | None = None, per_structure_targets: Sequence[str] = (), trainer_hypers: dict | None = None, checkpoint_dir: str = '') → None[source]

Train the scaler from data or load pre-trained weights.

This is the single source of truth for how to set up a scaler for use for preprocessing by any architecture.

Parameters:

  • scaler (Scaler) – The scaler to train or load into

  • train_datasets (List[Dataset | Subset]) – Training datasets

  • additive_models (List[Module]) – Additive models to subtract before fitting

  • batch_size (int) – Batch size for data loading

  • is_distributed (bool) – Whether training is distributed

  • fixed_weights (Dict[str, float | Dict[int, float]] | str | None) – Fixed weights dict, or path to a checkpoint. If it is a dictionary, the keys of the dict are the target names, and the values are either a single float value to be applied to all atomic types, or a dict mapping atomic type (int) to weight (float). If not provided, all scales will be computed based on the accumulated quantities. Fixed weights are not currently supported for atom-pair targets: their scales are always left at the identity value 1.0.

  • per_structure_targets (Sequence[str]) – Target names that should be treated as per-structure quantities and therefore not divided by the number of atoms.

  • trainer_hypers (dict | None) – Additional hyperparameters for the trainer.

  • checkpoint_dir (str) – Directory to save the scaler checkpoint

Return type:

None