Callback

class braintools.trainer.Callback

Base class for callbacks.

Callbacks allow you to hook into various stages of the training process. Override the methods you need in your custom callback subclass.

Examples

>>> class MyCallback(Callback):
...     def on_train_epoch_end(self, trainer, module):
...         print(f"Epoch {module.current_epoch} finished!")
...
>>> trainer = Trainer(callbacks=[MyCallback()])

load_state_dict(state_dict)

Load state from state dict.

on_after_backward(trainer, module)

Called after backward pass (gradient computation).

on_after_optimizer_step(trainer, module, optimizer)

Called after each optimizer step.

on_before_backward(trainer, module, loss)

Called before backward pass (gradient computation).

on_before_optimizer_step(trainer, module, optimizer)

Called before each optimizer step.

on_fit_end(trainer, module)

Called at the very end of fit.

on_fit_start(trainer, module)

Called at the very beginning of fit.

on_load_checkpoint(trainer, module, checkpoint)

Called when loading a checkpoint.

on_predict_batch_end(trainer, module, outputs, batch, batch_idx)

Called at the end of each predict batch.

on_predict_batch_start(trainer, module, batch, batch_idx)

Called at the beginning of each predict batch.

on_predict_end(trainer, module)

Called at the end of prediction.

on_predict_start(trainer, module)

Called at the beginning of prediction.

on_save_checkpoint(trainer, module, checkpoint)

Called when saving a checkpoint.

on_test_batch_end(trainer, module, outputs, batch, batch_idx)

Called at the end of each test batch.

on_test_batch_start(trainer, module, batch, batch_idx)

Called at the beginning of each test batch.

on_test_end(trainer, module)

Called at the end of testing.

on_test_epoch_end(trainer, module)

Called at the end of each test epoch.

on_test_epoch_start(trainer, module)

Called at the beginning of each test epoch.

on_test_start(trainer, module)

Called at the beginning of testing.

on_train_batch_end(trainer, module, outputs, batch, batch_idx)

Called at the end of each training batch.

on_train_batch_start(trainer, module, batch, batch_idx)

Called at the beginning of each training batch.

on_train_end(trainer, module)

Called at the end of training.

on_train_epoch_end(trainer, module)

Called at the end of each training epoch.

on_train_epoch_start(trainer, module)

Called at the beginning of each training epoch.

on_train_start(trainer, module)

Called at the beginning of training.

on_validation_batch_end(trainer, module, outputs, batch, batch_idx)

Called at the end of each validation batch.

on_validation_batch_start(trainer, module, batch, batch_idx)

Called at the beginning of each validation batch.

on_validation_end(trainer, module)

Called at the end of validation.

on_validation_epoch_end(trainer, module)

Called at the end of each validation epoch.

on_validation_epoch_start(trainer, module)

Called at the beginning of each validation epoch.

on_validation_start(trainer, module)

Called at the beginning of validation.

state_dict()

Return state dict for checkpointing.

Return type:

Dict[str, Any]

property state_key: str

Identifier for this callback in state dict.