Custom PyTorch-Lightning and Hydra Template for Deep Learning

---

Template Description (Delete this chapter when build your own project)

Efficient workflow and reproducibility are extremely important in every deep learning projects, so this template is designed with two great frameworks:

• Hydra for efficient configuration management and multiple jobs running. The key feature is the ability to dynamically create a hierarchical configuration by composition and override it through config files and the command line. All of that enables to easily switch between parameters and configurations without modifying the code.
• PyTorch Lightning, a PyTorch wrapper for fast prototyping and production. Please see PyTorchLightning Style Guide to understand the best practices with it.

Project Structure Overview

📦Custom-Lightning-Hydra-Template
 ┣ 📂checkpoints                            <- Checkpoints you may need in your training or inference
 ┣ 📂configs                                <- Root dir for Hydra
 ┃ ┣ 📂callbacks                            <- General callbacks configs, e.g., early stopping, model checkpointing, model summary and even progressbar.
 ┃ ┣ 📂datamodule                           <- PyTorchLightning Datamodule configs
 ┃ ┣ 📂experiment                           <- Experiment configs
 ┃ ┣ 📂extras                               <- Extras utilities configs
 ┃ ┣ 📂hydra                                <- Hydra configs
 ┃ ┣ 📂logger                               <- Logger configs
 ┃ ┣ 📂module                               <- PyTorchLightning Module configs
 ┃ ┣ 📂paths                                <- Project paths configs
 ┃ ┣ 📂trainer                              <- PyTorchLightning Trainer configs
 ┃ ┣ 📜inference.yaml                       <- Main config for inference
 ┃ ┗ 📜train.yaml                           <- Main config for training
 ┣ 📂data                                   <- Root dir for datasets
 ┣ 📂docs                                   <- Documentation for project
 ┣ 📂logs                                   <- Logs generated by Hydra and Lightning Loggers
 ┣ 📂notebooks                              <- Jupyter notebooks
 ┣ 📂scripts                                <- Useful shell scripts
 ┣ 📂src                                    <- Main source code
 ┃ ┣ 📂data                                 <- Dataset, Dataloader, DataModule, etc.
 ┃ ┃ ┗ 📂components                         <- Components, place Dataset class here if necessary
 ┃ ┣ 📂models                               <- PyTorchLightning Modules, place `LightningModule` here
 ┃ ┃ ┗ 📂components                         <- Components, place `torch.nn.Module` here
 ┃ ┣ 📂utils                                <- Utilities
 ┃ ┣ 📜inference.py                         <- Main script for inference
 ┃ ┗ 📜train.py                             <- Main script for training
 ┣ 📂tests                                  <- Tests
 ┣ 📜.gitignore                             <- List of files ignored by git
 ┣ 📜.gitmodules                            <- Submodules
 ┣ 📜.pre-commit-config.yaml                <- Pre-commit configuration for code formatting
 ┣ 📜.project-root                          <- For inferring the position of project root directory used by `rootutils`
 ┣ 📜README.md
 ┣ 📜environment.yaml
 ┗ 📜requirements.txt

How it works?

All PyTorch Lightning modules are dynamically instantiated from module paths specified in config. Example model config:

_target_: src.models.LitDDPM

ddpm:
  _target_: src.models.components.ddpm.DenoisingDiffusionProbabilisticModel
  timesteps: 1000
  beta_schedule: 'linear'
  unet:
    _target_: src.models.components.ddpm.UNetModel
    in_channels: 3
    model_channels: 128
    out_channels: 3

optimizer:
  _target_: torch.optim.Adam
  _partial_: true
  lr: 0.0002

compile: true

By using this config, we can instantiate your Lightning Module with Hydra (All related objects are instantiated in src/train.py):

model = hydra.utils.instantiate(config.model)

This allows you to easily iterate over new models! Every time you create a new one, just specify its module path and parameters in appropriate config file.

Switch between modules and datamodules with command line arguments using Hydra:

python src/train.py module=dcgan datamodule=celebahq

Which components are instantiated?

As you can see in main config configs/train.yaml which determines how config is composed when simply executing command python src/train.py, the following components are often instantiated in PyTorchLightning:

• data (Lightning DataModule)
• modules (Lightning Module)
• callbacks (Lightning/Custom Callbacks)
• logger (TensorBoard/Wandb/Comet/...)
• trainer (Lightning Trainer)

For different hyperparams, you can simply create a new experiment config which can overwrite the main config.

How to fit this template to my project?

We provide some examples, check them out in configs/experiment!

• ImageClassification:
  • ResNet + CIFAR10/Pneumonia
• ImageSegmentation:
  • UNet + BrainMRI-LGG
• StyleTransfer:
  • CycleGAN + Monet2Photo
• ImageGeneration:
  • Diffusion (UNet + DDPM) + CIFAR10/CelebA-HQ
  • Latent-Diffusion (VAE/VQGAN + UNet + DDPM/DDIM) + CelebA-HQ
• LanguageModeling:
  • GPT2(nanoGPT)

1. Define your `Dataset`, put it in `src/data/components/` if nec.

A Dataset defines how to process and access a single sample from your data. Any "Dataset" defined with __len__ and __getitem__ that can be used in an iterator DataLoader, is a valid Dataset, e.g., torch.utils.data.Dataset, hugginface's datasets.

Define transforms if nec., you can use torchvision and albumentations(RECOMMENDED) on vision tasks.

2. (Optional but Recommended) Fit your Dataset into `Lightning DataModule`, put it in `src/data/`.

As we usually use mini-batch training, data is organized in batched tensors by DataLoader. A DataLoader defines how to collate a batch of samples in a batched tensor. Any iterator defined with __len__ and __next__ that returns a batched tensor, is a valid DataLoader, e.g., PyTorch DataLoader, ffcv, NVIDIA DALI, etc.

Lightning DataModule helps you to split, transform, normalize your dataset, and defines a DataLoader used by Trainer.

Note that the core component in PytorchLightning is LightningModule and Trainer, so you can just pass a DataLoader to Trainer without using Lightning DataModule.

4. Fit your PyTorch Model (`nn.Module`) in `src/models/components` into Lightning Modules, put it in `src/models/`.

We provide a base class inherited from LightningModule and ABC named LitBaseModule in src/modules/base_module.py. A LitBaseModule have defined 2 additional hooks: get_input and get_output. The whole logic of the module is as following:

• __init__: get the configs of models, criterion, optimizer, lr_scheduler, etc., and instantiate them.
• get_input: parse the input batch into the model input.
• get_output: parse the model output into the final output.
• forward: the inference forward process of the module.
• setup: define the metric for different stages (fit, validate, test, predict).
• model_step: optional, the shared step which executes the model forward process and compute the loss.
• training_step: the training step of the module. NOTE that you should not do any logging in this function but in on_train_batch_end hook if you want to use torch.compile.
• other similar training/validation/test/predict hooks.
• configure_optimizers: choose what optimizers and learning-rate schedulers to use in your optimization.

See PyTorch Lightning Module for more details.

5. Write configs in `configs/`.

For simplicity, write `datamodule` and `module` configs first with your Lightning DataModule and Lightning Module, then create an `experiment` config to rewrite `datamodule`, `module`, `callbacks`, `trainer` and main configs.

6. Write tests in `tests/` and run them with `pytest`.

There're several things we need to check before running your deep learning experiment:

• Datamodule/dataloader (tests/test_datamodule.py). Check
  • if your datamodule can be instantiated
  • if the dataloader can be iterated
  • if the shape of data is correct
• Model (tests/test_model.py). Check
  • if your lightning module can be instantiated
  • if the i/o of the model (using torchinfo.summary) meets your design
• Misc (tests/test_configs.py). Check
  • if whole config (compose by hydra) is correct
  • if the three main components (datamodule, module, trainer) in main config (configs/train.yaml, configs/eval.yaml) can be instantiated
• Training (tests/test_train.py). Check
  • if training from scratch on CPU/GPU/16-mixed is correct
  • if resume from checkpoint is correct

For running:

# run all tests
pytest

# run w/ detail output
pytest -s

# run tests from specific file
pytest tests/test_train.py

# run tests from specific test
pytest tests/test_train.py::test_train_ddp_sim

# run all tests except the ones marked as slow
pytest -k "not slow"

6.5. (optional) Debug your module.

Sometimes, you may want to debug your module by running several steps/epochs/times of your module to check if it works as expected. (See https://lightning.ai/docs/pytorch/stable/debug/debugging_basic.html)

We've already provide config files in configs/debug/ and related option (debug) in train.yaml to help debug your module.

7. Run your experiment.

- Run training w/ chosen experiment config:

python src/train.py experiment=your_experiment_yaml_name

8. Export your module and inference.

We provide a base class inherited from `LitBaseModule` named `LitBaseModuleWithInference` in `src/modules/base_module.py` that defines functions to load ONNX and TensorRT models and do inference for single model (you might use multiple models in your inference pipeline). And you should implement your own inference pipeline in `inference_step` function.

For exporting your trained Lightning Module into ONNX and TensorRT, see src/export.py and configs/export.yaml. The basic logic is loading the trained checkpoint, exporting the model one by one (if you have multiple models in your pipeline.) and then doing some sanity check and pipeline test.

Advanced usage

1. Parallel training

In PytorchLightning, you can activate distributed training by setting args in Trainer (in this template, see configs/trainer).

There're several frequently used flags for distributed training:

• accelerator: Different accelerators (supports passing "cpu", "gpu", "tpu", "ipu", "auto")
  • The "auto" option recognizes the machine you are on, and selects the appropriate Accelerator.
• devices: Number of accelerators (devices=4)
• strategy: Supports passing different training strategy types ("ddp", "fsdp", "deepspeed", ...) as well as custom strategy instances.

For model parallelism, see https://lightning.ai/docs/pytorch/stable/advanced/model_parallel and https://lightning.ai/docs/pytorch/stable/advanced/model_init.html#model-parallel-training-fsdp-and-deepspeed

Besides, PL also supports use torchrun, see https://lightning.ai/docs/pytorch/stable/accelerators/gpu_intermediate.html.

2. Mixed-precision training

Lightning supports doing floating point operations in 64-bit precision (“double”), 32-bit precision (“full”), or 16-bit (“half”) with both regular and bfloat16. This selected precision will have a direct impact in the performance and memory usage based on your hardware. Automatic mixed precision settings are denoted by a "-mixed" suffix, while “true” precision settings have a "-true" suffix.

See PL's N-Bit Precision Tutorial for more details.

3. Hyperparameter optimization

4. Speed up your data loading

As we all know that LightningDataModule is just a wrapper of PyTorch DataLoader, so the speed of data loading is mainly affected by the speed of PyTorch DataLoader. On the other hand, LightningDataModule is only an option in PyTorch Lightning, so you can just pass something like a PyTorch DataLoader (essentially an iterator that outputs a batch of samples) to Trainer.

As alternatives, you can use DALI, FFCV, etc. to speed up your data loading. See https://lightning.ai/docs/pytorch/stable/data/alternatives.html for how to use them in PyTorch Lightning.

5. Automatic code formatting

Use pre-commit hooks to standardize code formatting of your project and save mental energy. Simply install pre-commit package with:

pip install pre-commit
pre-commit install
# optional
# pre-commit autoupdate
pre-commit run -a

6. Install project as a package

Supporting installing project as a package by changing the name of `src` folder to your project name and complete the `setup.py` file. Then install them, `pip install -e. / pip install git+git://github.com/YourGithubName/your-repo-name.git --upgrade`. So any file can be easily imported into any other file like:

from project_name.models.xx_module import LitxxxModule
from project_name.data.xxx_datamodule import XXXDataModule

7. Submodules related

First, you need to add a .gitmodules file in your project root directory by

touch .gitmodules

To add submodules into your project, run

git submodule add -b <branch> <repo-url> <local-path>

To install submodules if you forget to clone them when clone the main repo, using

git submodule update --init --recursive

To update submodules, using git submodule update --remote.

8. Deploying your module

To deploy our module, the first step is to load the checkpoint and the model config. You can load the checkpoint in PL and plain PyTorch.

Then you may need to export your model to onnx, torchscript, tensorrt, etc.

PL provides interfaces for you to export your model to ONNX and TorchScript.

For TensorRT, TODO.

Another thing you may need to do is to compress your model for fast inference w/ Quantization and Pruning.

See https://lightning.ai/docs/pytorch/stable/deploy/production.html

9. Profiling your module

If you want to find the bottlenecks in you code, see https://lightning.ai/docs/pytorch/stable/tuning/profiler.html and `configs/debug/profiler.yaml`.

10. Useful training tricks

PL provides many useful training tricks as `callbacks`, e.g., "Accumulate gradients", "Gradient clipping", "Stochastic Weight Averaging".

See https://lightning.ai/docs/pytorch/stable/advanced/training_tricks.html

11. Deploy your model

Currently, we support deploying your model to ONNX and TensorRT, see `configs/export.yaml` and `src/export.py` for more details.

To export your model to ONNX, there are two ways:

• Pure PyTorch (torch.onnx.export).
  1. Instantiate all the nn.Module model you need for inference and load the trained weights. You can also merge them to a single nn.Module and define __forward__ if necessary (That is what LightningModule does). DON'T forget to modify the layer name of checkpoints from PyTorchLightning.
  2. Define the input sample, input names, output names, dynamic axes, etc.
  3. Call torch.onnx.export to export the model.
• PyTorch Lightning (pl_module.to_onnx).
  1. Instantiate the LightningModule model and load the trained weights w/ pl_module.load_state_dict instead of LightningModuleClass.load_from_checkpoint.
  2. Same as pure PyTorch.
  3. Call pl_module.to_onnx to export the model.

Once you export your model to ONNX, you can simplify it using onnxsim.

Then, for inference w/ your ONNX model, install onnxruntime and use onnxruntime.InferenceSession to load your ONNX model and get predictions.

If you want to speedup the inference, you can using transfer your ONNX model into TensorRT engine. We provide a simple usage w/ TensorRT 10.1 in pure PyThon here (10.2 is buggy now).

1. Install TensorRT python package and Official CUDA Python (pycuda is a 3rd-party package, and is not recommended to use w/ newer TensorRT. E.g., pycuda is not work w/ TensorRT 10.x).
2. Build TensorRT engine from ONNX file w/ Python API. You can also use the binary tool trtexec provided by NVIDIA.
3. Load the deserialized TensorRT engine and get predictions, see the great tutorial and some supplementary use case for more details.

• fatemebafghi/TensorRT-Conversion-Calibration-and-inference
• TensorRT踩坑教程 - 银河渡舟

Supplymentary:

• ONNX Runtime Docs
• Export a PyTorch model to ONNX — PyTorch Tutorials documentation
• NVIDIA TensorRT Docs
• NVIDIA-AI-IOT/torch2trt: An easy to use PyTorch to TensorRT converter
• Convert PyTorch model to TensorRT for 3-8x speedup

Last but not least

Q1: How to customize the icon like ?

A1: See https://shields.io/docs/logos for more details.

---

Q2: How to customize the counter?

A2: See https://counter.seku.su/ or just search with keyword Moe-Counter.

Reference

• ashleve/lightning-hydra-template: PyTorch Lightning + Hydra. A very user-friendly template for ML experimentation. ⚡🔥⚡
• gorodnitskiy/yet-another-lightning-hydra-template: Flexible and scalable template based on PyTorch Lightning + Hydra. Efficient workflow and reproducibility for rapid ML experiments.
• Glossary — PyTorch Lightning 2.3.1 documentation
• How-to Guides — PyTorch Lightning 2.2.5 documentation

Recommendation toolkits

• albumentations: A fast and flexible image augmentation library for deep learning.
• timm: PyTorch image models, scripts and data.
• smp: Segmentation models with pretrained backbones. PyTorch.
• ultralytics: Ultralytics YOLO models.
• transformers: State-of-the-art Natural Language Processing for Pytorch and TensorFlow 2.0.
• Effective Training Techniques — PyTorch Lightning documentation

---

Description

What does your repo do?

Installation

# clone project
git clone https://github.com/YourGithubName/your-repo-name
cd your-repo-name

# [OPTIONAL] create conda environment
conda create -n yourenv python=3.x -y
conda activate yourenv

# install pytorch according to instructions
# https://pytorch.org/get-started/
conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia -y

# install requirements
pip install -r requirements.txt

# or you can install with `environment.yaml`

How to run

Data Preparation

Data preprocessing with progress bar (rich) and multiprocessing (ray).

# some instructions

Train

Train model with default configuration by choosing and modifying experiment configuration from configs/experiment/

# train SyncNet
python src/train.py trainer=ddp trainer.devices=2

You can override any parameter from command line like this

python src/train.py datamodule.batch_size=64

Inference

python src/inference.py # w/ some parameters

Reference