This is a project template for ML researchers developed at Mila. Our goal with this is to help you get started with a new research project.

See this introduction to the project for a detailed description of the context and motivations behind this project.

Please note: This is a Work-in-Progress. The goal is to make a first release by the end of summer 2024.

Installation

Projects created with this template use uv to manage dependencies. Once you have uv installed locally, you can install all dependencies with a single command:

uv sync  # Creates a virtual environment and installs dependencies in it.

For more detailed instructions, take a look at this page of the template docs.

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.

Who is this for? Why should you use this template?

This template comes with some unique features that can only be found here (for now)!

• Torch and Jax support
• Rich IDE support for Hydra config files
• Built-in automated tests (including reproducibility tests), including testing on SLURM clusters!
• And more! (see this page)

To make the best use of this template, you should ideally already have a good understanding of Python, some experience with PyTorch, and some basic experience with SLURM.

See this page for a list of other templates to choose from if this isn't for you.

Please consider making an issue on this repo if you feel like this could be improved, or something is confusing to you. We very much need and appreciate your feedback! 😊

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