Trainer

class lib.training.train.Trainer(plugin: TrainerBase, preview: bool, warmup_steps: int = 0, timelapse_folders: list[str] | None = None, timelapse_output: str = '')

Bases: object

Handles the feeding of training images to Faceswap models, the generation of Tensorboard logs and the creation of sample/time-lapse preview images.

All Trainer plugins must inherit from this class.

Parameters:

  • plugin (TrainerBase) – The plugin that will be processing each batch

  • preview (bool) – True to generate previews

  • warmup_steps (int) – The number of steps to warmup the learning rate for. Default: 0

  • timelapse_folders (list[str] | None) – The input folders to create timelapse images from. Default: None (no timelapse)

  • timelapse_output (str) – The folder to output timelapse images. Default: “” (no timelapse)

Attributes Summary

exit_early

True if the trainer should exit early, without performing any training steps

Methods Summary

save([is_exit])

Save the model

toggle_mask()

Toggle the mask overlay on or off based on user input.

train_one_batch()

Process a single batch through the model and obtain the loss

train_one_step(viewer[, do_timelapse])

Running training on a batch of images for each side.

Attributes Documentation

exit_early

True if the trainer should exit early, without performing any training steps

Methods Documentation

save(is_exit: bool = False) → None

Save the model

Parameters:

is_exit (bool) – True if save has been called on model exit. Default: False

Return type:

None

toggle_mask() → None

Toggle the mask overlay on or off based on user input.

Return type:

None

train_one_batch() → list[BatchLoss]

Process a single batch through the model and obtain the loss

Return type:

The collated loss values detached and moved to CPU in order (A, B, …)

train_one_step(viewer: Callable[[np.ndarray, str], None] | None, do_timelapse: bool = False) → None

Running training on a batch of images for each side.

Triggered from the training cycle in scripts.train.Train.

  • Runs a training batch through the model.

  • Outputs the iteration’s loss values to the console

  • Logs loss to Tensorboard, if logging is requested.

  • If a preview or time-lapse has been requested, then pushes sample images through the model to generate the previews

  • Creates a snapshot if the total iterations trained so far meet the requested snapshot criteria

Notes

As every iteration is called explicitly, the Parameters defined should always be None except on save iterations.

Parameters:

  • viewer (Callable[[np.ndarray, str], None] | None) – The function that will display the preview image

  • do_timelapse (bool) – True to generate a timelapse preview image

Return type:

None

property exit_early: bool

True if the trainer should exit early, without performing any training steps

save(is_exit: bool = False) → None

Save the model

Parameters:

is_exit (bool) – True if save has been called on model exit. Default: False

Return type:

None

toggle_mask() → None

Toggle the mask overlay on or off based on user input.

Return type:

None

train_one_batch() → list[BatchLoss]

Process a single batch through the model and obtain the loss

Return type:

The collated loss values detached and moved to CPU in order (A, B, …)

train_one_step(viewer: Callable[[np.ndarray, str], None] | None, do_timelapse: bool = False) → None

Running training on a batch of images for each side.

Triggered from the training cycle in scripts.train.Train.

  • Runs a training batch through the model.

  • Outputs the iteration’s loss values to the console

  • Logs loss to Tensorboard, if logging is requested.

  • If a preview or time-lapse has been requested, then pushes sample images through the model to generate the previews

  • Creates a snapshot if the total iterations trained so far meet the requested snapshot criteria

Notes

As every iteration is called explicitly, the Parameters defined should always be None except on save iterations.

Parameters:

  • viewer (Callable[[np.ndarray, str], None] | None) – The function that will display the preview image

  • do_timelapse (bool) – True to generate a timelapse preview image

Return type:

None