Add custom Train loop¶
VISSL implements a default training loop (single iteration step) that is used for self-supervised training of all VISSL reference approaches, for feature extraction and for supervised workflows. Users can implement their own training loop.
The training loop performs: data read, forward, loss computation, backward, optimizer step, parameter updates.
Various intermediate steps are also performed:
logging the training loss, training eta, LR, etc to loggers.
logging metrics to tensorboard.
performing any self-supervised method specific operations (like in MoCo approach, the momentum encoder is updated), computing the scores in swav.
checkpointing model if user wants to checkpoint in the middle of an epoch.
Users can implement their custom training loop by following the steps:
Step1: Create your
my_new_training_loop
module undervissl/trainer/train_steps/my_new_training_loop.py
following the template:
from vissl.trainer.train_steps import register_train_step
@register_train_step("my_new_training_loop")
def my_new_training_loop(task):
"""
add documentation on what this training loop does and how it varies from
standard training loop in vissl.
"""
# implement the training loop. It should take care of running the dataloader
# iterator to get the input sample
...
...
return task
Step2: New train loop is ready to use. Set the
TRAINER.TRAIN_STEP_NAME=my_new_training_loop