This is an alpha release. The API is subject to change and the documentation is incomplete.
The cortex
library provides a modular framework for neural network model composition for advanced use cases like multitask models, guided generative models, and multi-modal models.
Rather than tack on auxiliary abstractions to a single input --> single task model, cortex
is built from the ground up to make adding new inputs or tasks to the model as seamless as possible.
conda create --name cortex-env python=3.10 -y && conda activate cortex-env
python -m pip install pytorch-cortex
If you have a package version issue we provide pinned versions of all dependencies in requirements.txt
.
To update the frozen dependencies run
pip-compile --resolver=backtracking requirements.in
Use cortex_train_model --config-name <CONFIG_NAME>
to train, e.g.:
cortex_train_model --config-name train_protein_model wandb_mode=offline
- Configure the sweep
.yaml
, e.g../config/wandb/protein_model_sweep.yaml
- Run
wandb sweep ./config/wandb/protein_model_sweep.yaml
- Launch the wandb agents using a scheduler of your choice, e.g. SLURM or LSF
Contributions are welcome!
python -m pip install -r requirements-dev.in
pre-commit install
python -m pytest -v --cov-report term-missing --cov=./cortex ./tests
make -C docs html
cd docs/build/html
python -m http.server
Then open http://localhost:8000
in your browser.
A standalone whitepaper is in progress, in the meantime please cite the following paper if you use cortex
in your research:
@article{gruver2024protein,
title={Protein design with guided discrete diffusion},
author={Gruver, Nate and Stanton, Samuel and Frey, Nathan and Rudner, Tim GJ and Hotzel, Isidro and Lafrance-Vanasse, Julien and Rajpal, Arvind and Cho, Kyunghyun and Wilson, Andrew G},
journal={Advances in Neural Information Processing Systems},
volume={36},
year={2024}
}
The cerebral cortex is the outer layer of the brain (i.e. the grey matter). Different parts of the cortex are responsible for different functions, including sensation, motor function, and reasoning. The cortex is supported and interconnected by subcortical structures (i.e. the white matter). Similarly the cortex
package provides modular neural network components for different functions, such as extracting features from different types of inputs, or making predictions for different tasks.
cortex
also provides a flexible superstructure to compose those modules into advanced architectures and route signals between them.