graphium 2.2.0

Graphium: Scaling molecular GNNs to infinity.

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

mkdir ~/.venv                               # Create the folder for the environment
python3 -m venv ~/.venv/graphium_ipu        # Create the environment
source ~/.venv/graphium_ipu/bin/activate    # Activate the environment

# Install the PopTorch wheel
pip install PATH_TO_SDK/poptorch-3.2.0+109946_bb50ce43ab_ubuntu_20_04-cp38-cp38-linux_x86_64.whl

# Enable Poplar SDK (including Poplar and PopART)
source PATH_TO_SDK/enable

# Install the IPU specific and graphium requirements
pip install -r requirements_ipu.txt

# Install Graphium in dev mode
pip install --no-deps -e .

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

python main_run_multitask.py 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

python main_run_multitask.py 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

python main_run_multitask.py dataset=toymix model=gcn accelerator=gpu

automatically selects the correct configs to run the experiment on GPU. To use a config file you built from scratch you can run

python main_run_multitask.py --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