Welcome to repository and documentation the TrainTrack library. Detailed documentation coming very soon! See here for the documentation of the examples of this library.
TrainTrack is most easily installed with pip:
pip install traintrack
The aim of TrainTrack is simple: Given any set of self-contained Pytorch Lightning modules, run them in a serial and trackable way.
At its heart, TrainTrack is nothing more than a loop over the stages defined in a pipeline.yaml configuration file. However, it can also handle data processing steps (i.e. non-trainable modules), automatically creates grid scans over combinations of hyperparameters, logs training with (currently) either Tensorboard or Weights & Biases, and can run separate, dependent Slurm batch jobs. It also has an opinionated approach to how data is passed from stage to stage, via Lightning callbacks. In this way, the only code that needs to be written is Lightning modules, all other boilerplate and tracking is handled by TrainTrack.
traintrack uses two ingredients to run and track your training pipeline:
- A project configuration file
- A pipeline configuration file
It also makes one or two assumptions about the structure of your project. For project MyFirstMNIST, we should structure it as
📦 MyFirstMNIST
┣ 📂 architectures
┣ 📂 notebooks
┣ 📂 configs
┃ ┣ 📜 project_config.yaml
┃ ┗ 📜 my_first_pipeline.yaml
┗ 📂 logs
Note: Only configs/project_config.yaml is a required file. All else is configurable. An example project_config.yaml:
# project_config.yaml
# Location of libraries
libraries:
model_library: architectures
artifact_library: /my/checkpoint/directory
# The lines you would like/need in a batch script before the call to pipeline.py
custom_batch_setup:
- conda activate my-favorite-environment
# If you need to set up some environment before a batch is submitted, define it here in order of commands to run
command_line_setup:
- module load cuda
# If you need to run jobs serially, set to true
serial: False
# Which logger to use - options are Weights & Biases [wandb], TensorBoard [tb], or [None]
logger: wandb
We can launch a vanilla run of TrainTrack with
traintrack configs/my_first_pipeline.yaml
This trains and performs inference callbacks in the terminal.
The pipeline config file defines a pipeline, for example:
# my_first_pipeline.yaml
stages:
- {set: CNN, name: ResNet50, config: test_train.yaml}
which presumes a directory structure of:
📦 MyFirstMNIST
┣ 📂 architectures
┃ ┗ 📂 CNN
┃ ┃ ┣ 📜 cnn_base.py
┃ ┃ ┣ 📜 test_train.yaml
┃ ┃ ┗ 📂 Models
┃ ┃ ┃ ┗ 📜 resnet.py
Again, see this repository for example pipelines in action.