CNeuroMax is a Machine Learning workspace for model fitting (Deep Learning + Neuroevolution + HPO w/ Oríon), testing and serving (with Lightning Apps) AI/ML models. CNeuroMax aims to:
1. Reduce code & configuration boilerplate with:
- Hydra for task/experiment configuration.
- Hydra-zen for Hydra structured configuration management.
- Lightning for PyTorch code.
2. Simplify machine learning workflows:
- Hyperparameter optimization with Orion, AutoRL-Sweepers & Optuna
- SLURM job definition, queuing and monitoring with Submitit through its Hydra Launcher plugin.
- Docker / Apptainer environment containerization for both regular & SLURM-based execution.
- Transition from regular execution to SLURM-based execution by only swapping container technology and as little as a single Hydra configuration field.
3. Automate workspace & coding processes:
- Package upgrades through Renovate.
- Docstring documentation generation with Sphinx.
- Pre-commit formatting & linting hooks with pre-commit.
- Documentation/Docker image validation/deployment, formatting, linting,
type-checking & unit tests upon contribution to the
mainbranch using GitHub Actions.
4. Facilitate researcher collaboration through:
-
An object-oriented structure for code sharing & reusability.
-
A mono-repository workspace with task/experiment-specific subdirectories.
-
A very informative & clear to navigate Python API reference thanks to Autodoc and plugins like sphinx-autodoc-typehints and sphinx-paramlinks.
-
Shared logging with a Weights & Biases team space.
5. Promote high-quality and reproducible code by:
- Linting with Ruff, formatting with Black, unit-testing with pytest.
- Type-checking with Mypy (static) & Beartype (dynamic).
- DType & Shape type hinting for PyTorch tensors using jaxtyping & NumPy arrays using nptyping. Fully type checkable at runtime with Beartype.
- Providing a common Development Container recipe with the above features enabled + documentation preview with esbonio & GitHub Copilot.
6. Smoothen up rough edges by providing:
- Extensive documentation on how to install/execute on regular & SLURM-based systems.
- Unassuming guides on how to contribute to the codebase.
- Tutorials on i) how to facilitate code transport across machines & ii) how to prune unrelated components of the library for paper publication.
cneuromax/
├─ .github/ <-- Config files for GitHub Actions (tests, containers, etc)
├─ cneuromax/ <-- Root
│ ├─ fitting/ <-- Model fitting code
│ │ ├─ deeplearning/ <-- Deep Learning code
│ │ │ ├─ datamodule/ <-- Lightning DataModules
│ │ │ ├─ litmodule/ <-- Lightning Modules
│ │ │ │ └─ nnmodule/ <-- PyTorch Modules
│ │ │ ├─ utils/ <-- Deep Learning utilities
│ │ │ ├─ config.py <-- Deep Learning structured configs
│ │ │ ├─ runner.py <-- Deep Learning task runner
│ │ │ ├─ store.py <-- Deep Learning configs storing
│ │ │ └─ train.py <-- Deep Learning training function
│ │ ├─ deeplneuroevo/ <-- Deep Learning + Neuroevolution code
│ │ ├─ neuroevolution/ <-- Neuroevolution code
│ │ │ ├─ agent/ <-- Neuroevolution agents (encapsulate networks)
│ │ │ ├─ net/ <-- Neuroevolution networks
│ │ │ ├─ space/ <-- Neuroevolution spaces (where agents get evaluated)
│ │ │ ├─ utils/ <-- Neuroevolution utilities
│ │ │ ├─ config.py <-- Neuroevolution structured configs
│ │ │ ├─ evolve.py <-- Neuroevolution evolution function
│ │ │ └─ runner.py <-- Neuroevolution task runner
│ │ ├─ utils/ <-- Fitting utililities
│ │ ├─ config.py <-- Fitting structured configs
│ │ ├─ runner.py <-- Fitting task runner
│ │ └─ store.py <-- Fitting configs storing
│ ├─ projects/ <-- Contains all existing projects
│ │ │
│ │ │ ******************************************
│ │ └─ my_new_project/ <-- ******** Your new project folder *********
│ │ ├─ task/ <-- *********** Your task folder *************
│ │ │ └─ config.yaml <-- ****** Your task configuration file ******
│ │ ├─ datamodule.py <-- ******* Your Lightning DataModule ********
│ │ ├─ litmodule.py <-- ********* Your Lightning Module **********
│ │ └─ nnmodule.py <-- ********** Your PyTorch Module ***********
│ │ ******************************************
│ │
│ ├─ serving/ <-- Contains the code to create apps (cozmo inference, etc)
│ ├─ testing/ <-- Contains code for more complex testing of models
│ ├─ utils/ <-- CNeuroMax utilities
│ ├─ __init__.py <-- Sets up Beartype
│ ├─ __main__.py <-- Entry-point
│ ├─ config.py <-- Base structured configs
│ ├─ runner.py <-- Base task runner
│ └─ store.py <-- Base configs storing
├─ docs/ <-- Documentation files
├─ .devcontainer.json <-- VSCode container development config
├─ .gitignore <-- Files to not track with Git/GitHub
├─ .pre-commit-config.yaml <-- Pre-"git commit" actions config (format, lint, etc)
├─ .yamllint.yaml <-- YAML files config
├─ Dockerfile <-- To build the Docker image
├─ LICENSE <-- MIT License file
├─ README.md <-- Repository description file
├─ pyproject.toml <-- Python code & dependencies config
└─ renovate.json <-- Renovate Bot config (keeps dependencies up-to-date)
CNeuroMax is developed in the context of the Courtois Project on Neuronal Modelling, also known as CNeuroMod. Launched in 2018, CNeuroMod aims to create more human-like AI models by training them to emulate both human brain dynamics and behaviour. Phase I (2018-2023) of the project saw large-scale data acquisition and preliminary unimodal modelling. Phase II (2023-2027) of the project aims to create multimodal phantom models of the subjects. CNeuroMax is used as the framework to train these phantom models.