Please note: This is a Work-in-Progress. The goal is to make a first release by the end of summer 2024.
This is a template repository for a research project in machine learning. It is meant to be a starting point for new ML researchers that run jobs on SLURM clusters. The main target audience is Mila researchers and students, but this should still be useful to anyone that uses PyTorch-Lightning with Hydra.
For more context, see this introduction to the project..
This project makes use of the following libraries:
- Hydra is used to configure the project. It allows you to define configuration files and override them from the command line.
- PyTorch Lightning is used to as the training framework. It provides a high-level interface to organize ML research code.
- 🔥 Please note: You can also use Jax with this repo, as is shown in the Jax example 🔥
- Weights & Biases is used to log metrics and visualize results.
- pytest is used for testing.
Why should you use this template (instead of another)?
Here are some of the advantages to using this template compared to some of the other templates out there:
- ❗Can use both PyTorch and Jax with PyTorch-Lightning ❗
- Neat Configs thanks to Auto-Generated YAML schemas
- Easy development inside a Development Container with VsCode
- Tailor-made for ML researchers that run their jobs on SLURM clusters (with default configurations for the Mila and DRAC clusters.)
- Rich typing and documentation of all parts of the source code using Python 3.12's new type annotation syntax
- A comprehensive suite of automated tests for all algorithms, datasets and networks that are easy to reuse and extend
- Automatically creates Yaml Schemas for your Hydra config files (as soon as #7 is merged)
To see all available options:
python project/main.py --help
For a detailed list of examples, see the examples page.
pyproject.toml # Project metadata and dependencies
project/
main.py # main entry-point
algorithms/ # learning algorithms
datamodules/ # datasets, processing and loading
networks/ # Neural networks used by algorithms
configs/ # configuration files
docs/ # documentation
conftest.py # Test fixtures and utilities