Training

aido.training.pre_train(model: Surrogate, dataset: SurrogateDataset, n_epochs: int) → list[float]

Pre-train the Surrogate Model using a three-stage process.

This function performs pre-training in three stages with different batch sizes and learning rates to ensure stable convergence.

Parameters:

• model (Surrogate) – The surrogate model to pre-train.

• dataset (SurrogateDataset) – The dataset to use for training.

• n_epochs (int) – Number of epochs to train in each stage.

aido.training.training_loop(reco_file_paths_dict: dict | str | PathLike, reconstruction_loss_function: Callable[[Tensor, Tensor], Tensor], constraints: None | Callable[[SimulationParameterDictionary], float | Tensor] = None, wandb_logger: WandbLogger | None = None)

Internal training of the Surrogate and Optimizer models

Parameters:

• reco_file_paths_dict (dict | str | os.PathLike) – Either the dict with all the file paths or a single filepath (str or os.PathLike) that we first have to read from JSON.

• reconstruction_loss_function (Callable) – The user-defined loss function that provides the goodness of a given design. Has to take two Tensors (truth and predicted) and return a scalar Tensor used as the Optimizer loss.

• constraints (Callable, optional) – loss function, for example to account for cost penalties. Default is None

Returns:

SimulationParameterDictionary – The updated values as proposed by the Optimizer model.

Note

This function is integral to the correct training of the surrogate and optimizer models. The training itself consists of these steps: 1. Track all the file paths needed 2. Instantiate the Surrogate model if not done so, load it from .pt file if available from

current iteration (if training was stopped), then train it.

3. Run the Optimizer

4. Save results