train_les

train_les#

Classes#

Objective

Optuna objective class for hyperparameter optimization.

Functions#

load_dataset_and_sampler(conf, param_interior, ...)

Load the Z-score only dataset and sampler for training or validation.

load_model_states_and_optimizer(conf, model, device)

Load the model states, optimizer, scheduler, and gradient scaler.

main(rank, world_size, conf, backend[, trial])

Main function to set up training and validation processes.

primary_main()

Module Contents#

train_les.load_dataset_and_sampler(conf, param_interior, world_size, rank, is_train)#

Load the Z-score only dataset and sampler for training or validation.

train_les.load_model_states_and_optimizer(conf, model, device)#

Load the model states, optimizer, scheduler, and gradient scaler.

Parameters:

  • conf (dict) – Configuration dictionary containing training parameters.

  • model (torch.nn.Module) – The model to be trained.

  • device (torch.device) – The device (CPU or GPU) where the model is located.

Returns:

A tuple containing the updated configuration, model, optimizer, scheduler, and scaler.

Return type:

tuple

train_les.main(rank, world_size, conf, backend, trial=False)#

Main function to set up training and validation processes.

Parameters:

  • rank (int) – Rank of the current process.

  • world_size (int) – Number of processes participating in the job.

  • conf (dict) – Configuration dictionary containing model, data, and training parameters.

  • backend (str) – Backend to be used for distributed training.

  • trial (bool, optional) – Flag for whether this is an Optuna trial. Defaults to False.

Returns:

The result of the training process.

Return type:

Any

class train_les.Objective(config, metric='val_loss', device='cpu')#

Bases: echo.src.base_objective.BaseObjective

Optuna objective class for hyperparameter optimization.

config#

Configuration dictionary containing training parameters.

Type:

dict

metric#

Metric to optimize, defaults to “val_loss”.

Type:

str

device#

Device for training, defaults to “cpu”.

Type:

str

train(trial, conf)#

Train the model using the given trial configuration.

Parameters:

  • trial (optuna.trial.Trial) – Optuna trial object.

  • conf (dict) – Configuration dictionary for the current trial.

Returns:

The result of the training process.

Return type:

Any

train_les.primary_main()#