helios.trainer

Classes

TrainingUnit	Defines the types of units for training steps.
TrainingState	The training state.
Trainer	Automates the training, validation, and testing code.

Functions

get_trainer_safe_types_for_load(...)	Return the list of safe types for loading needed by the trainer.
register_trainer_types_for_safe_load(→ None)	Register trainer types for safe loading.
find_last_checkpoint(→ pathlib.Path | None)	Find the last saved checkpoint (if available).

Module Contents

class helios.trainer.TrainingUnit

Bases: enum.Enum

Defines the types of units for training steps.

ITERATION = 0

EPOCH = 1

classmethod from_str(label: str) → TrainingUnit

Convert the given string to the corresponding enum value.

Must be one of “iteration” or “epoch”.

Parameters:

label – the label to convert.

Returns:

The corresponding value.

Raises:

ValueError – if the given value is not one of “iteration” or “epoch”.

class helios.trainer.TrainingState

The training state.

Parameters:

• current_iteration – the current iteration number. Note that this refers to the iteration in which gradients are updated. This may or may not be equal to the global_iteration count.

• global_iteration – the total iteration count.

• global_epoch – the total epoch count.

• validation_cycles – the number of validation cycles.

• dataset_iter – the current batch index of the dataset. This is reset every epoch.

• early_stop_count – the current number of validation cycles for early stop.

• average_iter_time – average time per iteration.

• running_iter – iteration count in the current validation cycle. Useful for computing running averages of loss functions.

current_iteration: int = 0

global_iteration: int = 0

global_epoch: int = 0

validation_cycles: int = 0

dataset_iter: int = 0

early_stop_count: int = 0

average_iter_time: float = 0

running_iter: int = 0

dict

helios.trainer.get_trainer_safe_types_for_load() → list[Callable | tuple[Callable, str]]

Return the list of safe types for loading needed by the trainer.

Returns:

The list of types that need to be registered for safe loading.

helios.trainer.register_trainer_types_for_safe_load() → None

Register trainer types for safe loading.

helios.trainer.find_last_checkpoint(root: pathlib.Path | None) → pathlib.Path | None

Find the last saved checkpoint (if available).

The function assumes that checkpoint names contain epoch_<epoch> and iter_<iter> in the name, in which case it will return the path to the checkpoint with the highest epoch and/or iteration count.

Parameters:

root – the path where the checkpoints are stored.

Returns:

The path to the last checkpoint (if any).

class helios.trainer.Trainer(run_name: str = '', train_unit: TrainingUnit | str = TrainingUnit.EPOCH, total_steps: int | float = 0, valid_frequency: int | None = None, chkpt_frequency: int | None = None, print_frequency: int | None = None, accumulation_steps: int = 1, enable_cudnn_benchmark: bool = False, enable_deterministic: bool = False, early_stop_cycles: int | None = None, use_cpu: bool | None = None, gpus: list[int] | None = None, random_seed: int | None = None, enable_tensorboard: bool = False, enable_file_logging: bool = False, capture_warnings: bool = False, enable_progress_bar: bool = False, chkpt_root: pathlib.Path | None = None, log_root: pathlib.Path | None = None, wandb_args: helios.core.loggers.WandbArgs | None = None, src_root: pathlib.Path | None = None, import_prefix: str = '', print_banner: bool = True)

Automates the training, validation, and testing code.

The trainer handles all of the required steps to setup the correct environment (including handling distributed training), the training/validation/testing loops, and any clean up afterwards.

Parameters:

• run_name – (optional) name of the current run.

• train_unit – (optional) the unit used for training.

• total_steps – (optional) the total number of steps to train for.

• valid_frequency – (optional) frequency with which to perform validation.

• chkpt_frequency – (optional) frequency with which to save checkpoints.

• print_frequency – (optional) frequency with which to log.

• accumulation_steps – (optional) number of steps for gradient accumulation.

• enable_cudnn_benchmark – (optional) enable/disable CuDNN benchmark.

• enable_deterministic – (optional) enable/disable PyTorch deterministic.

• early_stop_cycles – (optional) number of cycles after which training will stop if no improvement is seen during validation.

• use_cpu – (optional) if true, CPU will be used.

• gpus – (optional) IDs of GPUs to use.

• random_seed – (optional) the seed to use for RNGs. Defaults to the value returned by get_default_seed().

• enable_tensorboard – (optional) enable/disable Tensorboard logging.

• enable_file_logging – (optional) enable/disable file logging.

• capture_warnings – (optional) enable/disable warnings capture in the log.

• enable_progress_bar – (optional) enable/disable the progress bar(s).

• chkpt_root – (optional) root folder in which checkpoints will be placed.

• log_root – (optional) root folder under which each active logger creates its own subfolder. Required when enable_file_logging, enable_tensorboard, or wandb_args is set.

• wandb_args – (optional) W&B configuration. Setting this enables W&B logging and requires log_root to be set. See WandbArgs for the available keys.

• src_root – (optional) root folder where the code is located. This is used to automatically populate the registries using update_all_registries().

• import_prefix – (optional) prefix to use when importing modules. See update_all_registries() for details.

• print_banner – (optional) if true, the Helios banner with system info will be printed.

property model: helios.model.Model

Return the model.

property datamodule: helios.data.DataModule

Return the datamodule.

property local_rank: int

Return the local rank of the trainer.

property rank: int

Return the global rank of the trainer.

property gpu_ids: list[int]

Return the list of GPU IDs to use for training.

property train_exceptions: list[type[Exception]]

Return the list of valid exceptions for training.

property test_exceptions: list[type[Exception]]

Return the list of valid exceptions for testing.

property plugins: dict[str, helios.plugins.Plugin]

Return the list of plug-ins.

register_plugin(plugin: helios.plugins.Plugin) → None

Register a plug-in with the trainer.

The plug-in is registered using its plug_id.

Note

configure_trainer() and configure_model() are called automatically by the trainer before training/testing.

Parameters:

plugin – the plug-in to register.

Raises:

• KeyError – if a plug-in with the same ID is already registered.

• ValueError – if the plug-in conflicts with an already-registered plug-in’s UniquePluginOverrides.

property queue: torch.multiprocessing.Queue | None

Return the multi-processing queue instance.

Note

If training isn’t distributed or if torchrun, then None is returned instead.

fit(model: helios.model.Model, datamodule: helios.data.DataModule) → bool

Run the full training routine.

Parameters:

• model – the model to run on.

• datamodule – the datamodule to use.

Returns:

True if the training process completed successfully, false otherwise.

test(model: helios.model.Model, datamodule: helios.data.DataModule) → bool

Run the full testing routine.

Parameters:

• model – the model to run on.

• datamodule – the datamodule to use.

Returns:

True if the training process completed successfully, false otherwise.