train_wrf_multi

train_wrf_multi#

Classes#

Objective

Optuna objective class for hyperparameter optimization.

Functions#

load_dataset_and_sampler(conf, param_interior, ...[, ...])

load_model_states_and_optimizer(conf, model, device)

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

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

primary_main()

Module Contents#

train_wrf_multi.load_dataset_and_sampler(conf, param_interior, param_outside, world_size, rank, is_train=True)#
train_wrf_multi.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_wrf_multi.main(rank, world_size, conf, backend, trial=False)#
class train_wrf_multi.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_wrf_multi.primary_main()#