towhee.trainer.callback.ModelCheckpointCallback

class towhee.trainer.callback.ModelCheckpointCallback(trainercontrol: TrainerControl, filepath: str = './', every_n_epoch: int = -1, every_n_iteration: int = -1)[source]

Bases: Callback

ModelCheckpointCallback is intended to save the model at some interval. It can be set in epoch mode or iteration mode. Only one of every_n_epoch and every_n_iteration can be set to a positive value and the trainer.should_save will set to True when the condion meets.

Parameters:

  • trainercontrol (TrainerControl) – The trainercontrol which callback can operate.

  • filepath (str) – Filepath to save the model.

  • every_n_epoch (int) – Save the model after n epochs.

  • every_n_iteration (int) – Save the model after n iterations.

Methods

on_batch_begin

Hook function invoked before every batch calculation.

on_batch_end

Hook function invoked after every batch calculation.

on_epoch_begin

Hook function invoked before each epoch.

on_epoch_end

Hook function invoked after each epoch.

on_eval_batch_begin

Hook function invoked before every batch calculation in evaluate stage.

on_eval_batch_end

Hook function invoked after every batch calculation in evaluate stage.

on_eval_begin

Hook function invoked before evaluate stage.

on_eval_end

Hook function invoked after evaluate stage.

on_train_batch_begin

Hook function invoked before train stage.

on_train_batch_end

Hook function invoked before every batch calculation in train stage.

on_train_begin

Hook function invoked before train stage.

on_train_end

Hook function invoked after train stage.

set_model

Set the model to callback.

set_optimizer

Set the optimizer to callback.

set_trainercontrol

Set the trainercontrol to callback.
__init__(trainercontrol: TrainerControl, filepath: str = './', every_n_epoch: int = -1, every_n_iteration: int = -1)[source]
on_batch_begin(batch: Tuple, logs: Dict) → None

Hook function invoked before every batch calculation.

Parameters:

  • batch (Tuple) – The data batch to calculate.

  • logs (Dict) – Kv store to save and load info.

on_batch_end(batch: Tuple, logs: Optional[Dict] = None)[source]

Hook function invoked after every batch calculation.

Parameters:

  • batch (Tuple) – The data batch to calculate.

  • logs (Dict) – Kv store to save and load info.

on_epoch_begin(epochs: int, logs: Dict) → None

Hook function invoked before each epoch.

Parameters:

  • epochs (int) – Epoch index.

  • logs (Dict) – Kv store to save and load info.

on_epoch_end(epochs: int, logs: Optional[Dict] = None)[source]

Hook function invoked after each epoch.

Parameters:

  • epochs (int) – Epoch index.

  • logs (Dict) – Kv store to save and load info.

on_eval_batch_begin(batch: Tuple, logs: Dict) → None

Hook function invoked before every batch calculation in evaluate stage.

Parameters:

  • batch (Tuple) – The data batch to calculate.

  • logs (Dict) – Kv store to save and load info.

on_eval_batch_end(batch: Tuple, logs: Dict) → None

Hook function invoked after every batch calculation in evaluate stage.

Parameters:

  • batch (Tuple) – The data batch to calculate.

  • logs (Dict) – Kv store to save and load info.

on_eval_begin(logs: Dict) → None

Hook function invoked before evaluate stage.

Parameters:

logs (Dict) – Kv store to save and load info.

on_eval_end(logs: Dict) → None

Hook function invoked after evaluate stage.

Parameters:

logs (Dict) – Kv store to save and load info.

on_train_batch_begin(batch: Tuple, logs: Dict) → None

Hook function invoked before train stage.

Parameters:

logs (Dict) – Kv store to save and load info.

on_train_batch_end(batch: Tuple, logs: Dict) → None

Hook function invoked before every batch calculation in train stage.

Parameters:

  • batch (Tuple) – The data batch to calculate.

  • logs (Dict) – Kv store to save and load info.

on_train_begin(logs: Dict) → None

Hook function invoked before train stage.

Parameters:

logs (Dict) – Kv store to save and load info.

on_train_end(logs: Dict) → None

Hook function invoked after train stage.

Parameters:

logs (Dict) – Kv store to save and load info.

set_model(model: Module) → None

Set the model to callback.

Parameters:

model (torch.nn.Module) – The model which callback can operate.

set_optimizer(optimizer: Optimizer) → None

Set the optimizer to callback.

Parameters:

optimizer (torch.optim.Optimizer) – The optimizer which callback can operate.

set_trainercontrol(trainercontrol: TrainerControl) → None

Set the trainercontrol to callback.

Parameters:

trainercontrol (towhee.trainer.callback.TrainerControl) – The trainercontrol which callback can operate.