Package: graphium
Scaling molecular GNNs to infinityA 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 onPytorch 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!DocumentationVisithttps://graphium-docs.datamol.io/.You can try running Graphium on Graphcore IPUs for free on Gradient by clicking on the button above.Installation for developersFor CPU and GPU developersUsemamba, a faster and better alternative toconda.If you are using a GPU, we recommend enforcing the CUDA version that you need withCONDA_OVERRIDE_CUDA=XX.X.# Install Graphium's dependencies in a new environment named `graphium`mambaenvcreate-fenv.yml-ngraphium# To force the CUDA version to 11.2, or any other version you prefer, use the following command:# CONDA_OVERRIDE_CUDA=11.2 mamba env create -f env.yml -n graphium# Install Graphium in dev modemambaactivategraphium
pipinstall--no-deps-e.For IPU developers# Install Graphcore's SDK and Graphium dependencies in a new environment called `.graphium_ipu`./install_ipu.sh.graphium_ipuThe above step needs to be done once. After that, enable the SDK and the environment as follows:sourceenable_ipu.sh.graphium_ipuTraining a modelTo learn how to train a model, we invite you to look at the documentation, or the jupyter notebooks availablehere.If you are not familiar withPyTorchorPyTorch-Lightning, we highly recommend going through their tutorial first.Running an experimentWe have setup Graphium withhydrafor managing config files. To run an experiment go to theexpts/folder. For example, to benchmark a GCN on the ToyMix dataset rungraphium-trainarchitecture=toymixtasks=toymixtraining=toymixmodel=gcnTo change parameters specific to this experiment like switching fromfp16tofp32precision, you can either override them directly in the CLI viagraphium-trainarchitecture=toymixtasks=toymixtraining=toymixmodel=gcntrainer.trainer.precision=32or change them permanently in the dedicated experiment config underexpts/hydra-configs/toymix_gcn.yaml.
Integratinghydraalso allows you to quickly switch between accelerators. E.g., runninggraphium-trainarchitecture=toymixtasks=toymixtraining=toymixmodel=gcnaccelerator=gpuautomatically selects the correct configs to run the experiment on GPU.
Finally, you can also run a fine-tuning loop:graphium-train+finetuning=admetTo use a config file you built from scratch you can rungraphium-train--config-path[PATH]--config-name[CONFIG]Thanks to the modular nature ofhydrayou can reuse many of our config settings for your own experiments with Graphium.Preparing the data in advanceThe data preparation including the featurization (e.g., of molecules from smiles to pyg-compatible format) is embedded in the pipeline and will be performed when executinggraphium-train [...].However, when working with larger datasets, it is recommended to perform data preparation in advance using a machine with sufficient allocated memory (e.g., ~400GB in the case ofLargeMix). Preparing data in advance is also beneficial when running lots of concurrent jobs with identical molecular featurization, so that resources aren't wasted and processes don't conflict reading/writing in the same directory.The following command-line will prepare the data and cache it, then use it to train a model.# First prepare the data and cache it in `path_to_cached_data`graphiumdataprepare++datamodule.args.processed_graph_data_path=[path_to_cached_data]# Then train the model on the prepared datagraphium-train[...]datamodule.args.processed_graph_data_path=[path_to_cached_data]Notethatdatamodule.args.processed_graph_data_pathcan also be specified atexpts/hydra_configs/.Notethat, every time the configs ofdatamodule.args.featurizationchanges, you will need to run a new data preparation, which will automatically be saved in a separate directory that uses a hash unique to the configs.LicenseUnder the Apache-2.0 license. SeeLICENSE.DocumentationDiagram for data processing in Graphium.Diagram for Muti-task network in Graphium