Research Project Template#

Work-in-Progress

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 research project template. It is meant to be a starting point for ML researchers at Mila.

For more context, see this introduction to the project..

Set up in 5 minutes

Get started quickly with a single installation script and get up and running in minutes

Getting started
Well-tested, robust codebase

Focus on your research! Let tests take care of detecting bugs and broken configs!

Check out the included tests
Support for both PyTorch and Jax

You can use both PyTorch and Jax for your algorithms! (Lightning handles the rest.)

Check out the Jax example
Ready-to-use examples

Includes examples for Supervised learning(1) and NLP 🤗, with unsupervised learning and RL coming soon.
1. The source code for the example is available here
Check out the examples here

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 described in the Jax example 🔥
Weights & Biases is used to log metrics and visualize results.
pytest is used for testing.

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