mml.core.scripts.callbacks

class MMLModelCheckpoint[source]

Bases: ModelCheckpoint

Slight modification of the lightning ModelCheckpoint (see https://github.com/Lightning-AI/pytorch-lightning/issues/20245).

on_train_epoch_end(trainer: Trainer, pl_module: LightningModule) → None[source]: Save a checkpoint at the end of the training epoch.

on_validation_end(trainer: Trainer, pl_module: LightningModule) → None[source]: Save a checkpoint at the end of the validation stage.

class MMLRichProgressBar[source]

Bases: RichProgressBar

Slight modification of the Lightning rich progress bar, showing the correct experiment name.

get_metrics(trainer, model)[source]

class MMLTQDMProgressBar[source]

Bases: TQDMProgressBar

Slight modification of the Lightning tqdm progress bar, showing the correct experiment name.

__init__(refresh_rate=1)[source]

get_metrics(trainer, model)[source]

class MetricsTrackerCallback[source]

Bases: Callback

Keeps track of all metrics, at the end of each epoch.

__init__()[source]

copy_metrics(trainer: Trainer, phase: LearningPhase) → None[source]

load_state_dict(state_dict: Dict[str, Any]) → None[source]

on_test_end(trainer: Trainer, pl_module: LightningModule) → None[source]

on_train_epoch_end(trainer: Trainer, pl_module: LightningModule) → None[source]

on_validation_epoch_end(trainer: Trainer, pl_module: LightningModule) → None[source]

state_dict() → Dict[str, Any][source]

class StopAfterKeyboardInterrupt[source]

Bases: Callback

Ensures pytorch lightning to really shut down after keyboard interrupt. This is the new variant of this Callback aimed to be used by most recent pytorch lightning versions.

on_exception(trainer: Trainer, pl_module: LightningModule, exception: BaseException) → None[source]