Research Project Template

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..

Overview

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 use this template?

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)

Usage

To see all available options:

python project/main.py --help

For a detailed list of examples, see the examples page.

Project layout

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