Graphium: Scaling molecular GNNs to infinity.
Project description
Scaling molecular GNNs to infinity
A deep learning library focused on graph representation learning for real-world chemical tasks.
- ✅ State-of-the-art GNN architectures.
- 🐍 Extensible API: build your own GNN model and train it with ease.
- ⚗️ Rich featurization: powerful and flexible built-in molecular featurization.
- 🧠 Pretrained models: for fast and easy inference or transfer learning.
- ⮔ Read-to-use training loop based on Pytorch Lightning.
- 🔌 Have a new dataset? Graphium provides a simple plug-and-play interface. Change the path, the name of the columns to predict, the atomic featurization, and you’re ready to play!
Documentation
Visit https://graphium-docs.datamol.io/.
You can try running Graphium on Graphcore IPUs for free on Gradient by clicking on the button above.
Installation for developers
For CPU and GPU developers
Use mamba
:
# Install Graphium's dependencies in a new environment named `graphium`
mamba env create -f env.yml -n graphium
# Install Graphium in dev mode
mamba activate graphium
pip install --no-deps -e .
For IPU developers
# Install Graphcore's SDK and Graphium dependencies in a new environment called `.graphium_ipu`
./install_ipu.sh .graphium_ipu
The above step needs to be done once. After that, enable the SDK and the environment as follows:
source enable_ipu.sh .graphium_ipu
Training a model
To learn how to train a model, we invite you to look at the documentation, or the jupyter notebooks available here.
If you are not familiar with PyTorch or PyTorch-Lightning, we highly recommend going through their tutorial first.
Running an experiment
We have setup Graphium with hydra
for managing config files. To run an experiment go to the expts/
folder. For example, to benchmark a GCN on the ToyMix dataset run
graphium-train dataset=toymix model=gcn
To change parameters specific to this experiment like switching from fp16
to fp32
precision, you can either override them directly in the CLI via
graphium-train dataset=toymix model=gcn trainer.trainer.precision=32
or change them permamently in the dedicated experiment config under expts/hydra-configs/toymix_gcn.yaml
.
Integrating hydra
also allows you to quickly switch between accelerators. E.g., running
graphium-train dataset=toymix model=gcn accelerator=gpu
automatically selects the correct configs to run the experiment on GPU. Finally, you can also run a fine-tuning loop:
graphium-train +finetuning=admet
To use a config file you built from scratch you can run
graphium-train --config-path [PATH] --config-name [CONFIG]
Thanks to the modular nature of hydra
you can reuse many of our config settings for your own experiments with Graphium.
License
Under the Apache-2.0 license. See LICENSE.
Documentation
- Diagram for data processing in molGPS.
- Diagram for Muti-task network in molGPS
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.