ModelTrainer Module

[gitttt-1234]

ModelTrainer Module is required to facilitate the training of different types of model using PyTorch Lightning with optional logging in wandb. We define the ModelTrainer Module as follows:

ModelTrainer Module: This class is used to train a PyTorch model using Lightning and save the trained model. The inputs to this class are config files for Data pipeline, model backbone and head configs and training configs. We have a child class of pl.LightningModule for each type of model - SingleInstance, TopDownCenteredInstance, etc, which is trained using a Trainer Class.

The ModelTrainer Module performs the following functions:

1. Creates DataLoaders from the data configs
2. Trains a Lightning Model with Trainer class
3. Optional Logging with wandb to track the training process
4. Saves Checkpoints throughout the training

from omegaconf import OmegaConf
from sleap_nn.models import ModelTrainer

config= OmegaConfig.create() # data config, model config and trainer config

m = ModelTrainer(config)
m.train() # calls the trainer.fit() from Lightning
# the checkpoints are saved in the specified directory

[gitttt-1234]

PR #29