API Reference

Trainer

AcceleratorModule

ExtendedAcceleratorModule

States

class accmt.states.TrainingState(batch_iteration: int = 0, global_step: int = 0, train_step: int = 0, val_step: int = 0, epoch: int = 0, is_end_of_epoch: bool = False, is_last_training_batch: bool = False, is_last_validation_batch: bool = False, is_last_epoch: bool = False, evaluations_done: int = 0, additional_metrics: dict[str, dict[str, ~typing.Any]] = <factory>, patience_left: dict[str, int] = <factory>, best_train_loss: float = inf, finished: bool = False, num_checkpoints_made: int = 0, train_dataloader_idx: int = 0)

General training state.

Parameters:

• batch_iteration (int) – Batch iteration index. This is incremented every time a train step is done. This ignores gradient accumulation. If gradient accumulation is not used, this is the same as global_step.

• global_step (int) – Global step index. This is incremented every time a train step and a gradient accumulation step is done.

• train_step (int) – Training step index inside a training dataloader (can be considered as batch index).

• val_step (int) – Validation step index inside an evaluation loop (can be considered as batch index).

• epoch (int) – Epoch index.

• is_end_of_epoch (bool) – Flag to check if current state is at the end of an epoch.

• is_last_training_batch (bool) – Flag to check if current state is processing the last training batch.

• is_last_validation_batch (bool) – Flag to check if current state is processing the last validation batch

• is_last_epoch (bool) – Flag to check if current state is processing the last epoch.

• evaluations_done (int) – Number of evaluations done.

• additional_metrics (dict) – Additional metrics (e.g. accuracy, bleu, f1, etc).

• model_savings (dict) – Bests model saving values.

• patience_left (dict) – Patience left per model saving (in case it’s implemented, otherwise values are set to -1).

• best_train_loss (float) – Best training loss achieved.

• best_valid_loss (float) – Best validation loss achieved.

• finished (bool, optional, defaults to False) – Flag to identify if the process has already finished.

• num_checkpoints_made (int, optional, defaults to 0) – Number of checkpoints made.

• train_dataloader_idx (int, optional, defaults to None) – Index of the train dataloader being used.

additional_metrics: dict[str, dict[str, Any]]

best_train_loss: float = inf

epoch: int = 0

evaluations_done: int = 0

finished: bool = False

global_step: int = 0

is_end_of_epoch: bool = False

is_last_epoch: bool = False

is_last_training_batch: bool = False

is_last_validation_batch: bool = False

num_checkpoints_made: int = 0

patience_left: dict[str, int]

train_step: int = 0

val_step: int = 0

Callbacks

class accmt.callbacks.Callback(module: AcceleratorModule = None, state: TrainingState = None)

Callback module containing different callback functions for different stages of the traininig process.

NOTE: Every callback function will run on every process. If you want your callback functions to only run on a single process, make sure to import accmt.decorators for different function decorators.

Variables:

• module (AcceleratorModule) – Training module.

• trainer (Trainer) – Defined Trainer class.

• state (TrainingState) – Reference to TrainingState class.

on_fit_start(*optional*)

Callback when training process starts.

on_fit_end(*optional*)

Callback when training process ends.

on_before_backward(*optional*)

Callback before engine’s backward.

on_after_backward(*optional*)

Callback after engine’s backward.

on_before_optimizer_step(*optional*)

Callback before optimizers steps.

on_after_optimizer_step(*optional*)

Callback after optimizer steps.

on_before_scheduler_step(*optional*)

Callback before scheduler steps:

on_after_scheduler_step(*optional*)

Callback after scheduler steps.

on_before_zero_grad(*optional*)

Callback before optimizer resets gradients.

on_after_zero_grad(*optional*)

Callback after optimizer resets gradients.

on_runtime_error(*optional*)

Callback when process raises a RunTimeError exception.

on_cuda_out_of_memory(*optional*)

Callback when process raises a RunTimeError exception with CUDA Out Of Memory.

on_keyboard_interrupt(*optional*)

Callback when process raises a KeyboardInterrupt exception.

on_exception(*optional*)

Callback when process raises any other Exception different than RuntimeError and KeyboardInterrupt

on_resume(*optional*)

Callback when resuming training process.

on_save_checkpoint(*optional*)

Callback when saving checkpoint.

on_before_training_step(*optional*)

Callback before training_step function.

on_after_training_step(*optional*)

Callback after training_step function.

on_before_validation_step(*optional*)

Callback before validation_step function.

on_after_validation_step(*optional*)

Callback after validation_step function.

on_epoch_start(*optional*)

Callback when an epoch starts.

on_epoch_end(*optional*)

Callback when an epoch ends.

on_evaluation_start(*optional*)

Callback when evaluation starts.

on_evaluation_end(*optional*)

Callback when evaluation ends.

module: AcceleratorModule = None

on_after_backward()

Callback after engine’s backward.

on_after_optimizer_step(optimizer: Optimizer)

Callback after optimizer steps.

Parameters:

optimizer (Optimizer) – Wrapped optimizer.

on_after_scheduler_step(scheduler: LRScheduler)

Callback after scheduler steps.

Parameters:

scheduler (LRScheduler) – Wrapped scheduler.

on_after_training_step()

Callback after training_step function.

on_after_validation_step()

Callback after validation_step function.

on_after_zero_grad(optimizer: Optimizer)

Callback after optimizer resets gradients.

Parameters:

optimizer (Optimizer) – Wrapped optimizer.

on_before_backward(loss: Tensor)

Callback before engine’s backward.

Parameters:

loss (torch.Tensor) – Scalar loss tensor.

on_before_optimizer_step(optimizer: Optimizer)

Callback before optimizers steps.

Parameters:

optimizer (Optimizer) – Wrapped optimizer.

on_before_scheduler_step(scheduler: LRScheduler)

Callback before scheduler steps:

Parameters:

scheduler (LRScheduler) – Wrapped scheduler.

on_before_training_step(batch: Any)

Callback before training_step function.

Parameters:

batch (Any) – Dataloader’s batch.

on_before_validation_step(batch: Any)

Callback before validation_step function.

Parameters:

batch (Any) – Dataloader’s batch.

on_before_zero_grad(optimizer: Optimizer)

Callback before optimizer resets gradients.

Parameters:

optimizer (Optimizer) – Wrapped optimizer.

on_cuda_out_of_memory(exception: Exception)

Callback when process raises a RunTimeError exception with CUDA Out Of Memory.

Parameters:

exception (Exception) – Raised exception.

on_epoch_end()

Callback when an epoch ends.

on_epoch_start()

Callback when an epoch starts.

on_evaluation_end()

Callback when evaluation ends.

on_evaluation_start()

Callback when evaluation starts.

on_exception(exception: Exception)

Callback when process raises any other Exception different than RuntimeError and KeyboardInterrupt

Parameters:

exception (Exception) – Raised exception.

on_fit_end()

Callback when training process ends.

on_fit_start()

Callback when training process starts.

on_keyboard_interrupt(exception: Exception)

Callback when process raises a KeyboardInterrupt exception.

Parameters:

exception (Exception) – Raised exception.

on_resume()

Callback when resuming training process.

on_runtime_error(exception: Exception)

Callback when process raises a RunTimeError exception.

Parameters:

exception (Exception) – Raised exception.

on_save_checkpoint()

Callback when saving checkpoint.

state: TrainingState = None

trainer = None

Metrics

DataCollators

Monitor

HyperParameters

Optimizers

Schedulers

HyperParameterSearch