Merge pull request #405 from datamol-io/hydra-integration

Hydra Integration

Author: DomInvivo

README.md

@@ -19,6 +19,7 @@
 [![code-check](https://github.com/datamol-io/graphium/actions/workflows/code-check.yml/badge.svg)](https://github.com/datamol-io/graphium/actions/workflows/code-check.yml)
 [![doc](https://github.com/datamol-io/graphium/actions/workflows/doc.yml/badge.svg)](https://github.com/datamol-io/graphium/actions/workflows/doc.yml)
 [![codecov](https://codecov.io/gh/datamol-io/graphium/branch/main/graph/badge.svg?token=bHOkKY5Fze)](https://codecov.io/gh/datamol-io/graphium)
+[![hydra](https://img.shields.io/badge/Config-Hydra_1.3-89b8cd)](https://hydra.cc/)
 
 A deep learning library focused on graph representation learning for real-world chemical tasks.
 
@@ -78,6 +79,27 @@ To learn how to train a model, we invite you to look at the documentation, or th
 
 If you are not familiar with [PyTorch](https://pytorch.org/docs) or [PyTorch-Lightning](https://pytorch-lightning.readthedocs.io/en/latest/), 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
+```bash
+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
+```bash
+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
+```bash
+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
+```bash
+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](LICENSE).

env.yml

@@ -71,3 +71,4 @@ dependencies:
 
   - pip:
       - lightning-graphcore # optional, for using IPUs only
+      - hydra-core>=1.3.2

expts/hydra-configs/accelerator/cpu.yaml

@@ -0,0 +1 @@
+type: cpu

expts/hydra-configs/accelerator/gpu.yaml

@@ -0,0 +1 @@
+type: gpu

expts/hydra-configs/accelerator/ipu.yaml

@@ -0,0 +1,9 @@
+type: ipu
+ipu_config:
+    - deviceIterations(5) # IPU would require large batches to be ready for the model.
+    - replicationFactor(16)
+    # - enableProfiling("graph_analyser")       # The folder where the profile will be stored
+    # - enableExecutableCaching("pop_compiler_cache")
+    - TensorLocations.numIOTiles(128)
+    - _Popart.set("defaultBufferingDepth", 128)
+    - Precision.enableStochasticRounding(True)

expts/hydra-configs/dataset/accelerator/toymix_cpu.yaml

@@ -0,0 +1,27 @@
+# @package _global_
+
+architecture:
+  task_heads:
+    tox21:
+      last_activation: none
+
+datamodule:
+  args:
+    batch_size_training: 200
+    batch_size_inference: 200
+    featurization_n_jobs: 4
+    num_workers: 4
+
+predictor:
+  optim_kwargs: {}
+  loss_fun:
+    tox21: bce_logits_ipu
+  metrics_every_n_steps: 300
+  torch_scheduler_kwargs:
+    max_num_epochs: &max_epochs 300
+
+trainer:
+  trainer:
+    precision: 32
+    accumulate_grad_batches: 1
+    max_epochs: *max_epochs

expts/hydra-configs/dataset/accelerator/toymix_gpu.yaml

@@ -0,0 +1,29 @@
+# @package _global_
+
+accelerator:
+  float32_matmul_precision: medium
+
+architecture:
+  task_heads:
+    tox21:
+      last_activation: none
+
+datamodule:
+  args:
+    batch_size_training: 200
+    batch_size_inference: 200
+    featurization_n_jobs: 4
+    num_workers: 4
+
+predictor:
+  optim_kwargs: {}
+  loss_fun:
+    tox21: bce_logits_ipu
+  metrics_every_n_steps: 300
+  torch_scheduler_kwargs:
+    max_num_epochs: &max_epochs 300
+
+trainer:
+  trainer:
+    accumulate_grad_batches: 1
+    max_epochs: *max_epochs

expts/hydra-configs/dataset/accelerator/toymix_ipu.yaml

@@ -0,0 +1,23 @@
+# @package _global_
+
+datamodule:
+  args:
+    ipu_dataloader_training_opts:
+      mode: async
+      max_num_nodes_per_graph: 44 # train max nodes: 20, max_edges: 54
+      max_num_edges_per_graph: 80
+    ipu_dataloader_inference_opts:
+      mode: async
+      max_num_nodes_per_graph: 44 # valid max nodes: 51, max_edges: 118
+      max_num_edges_per_graph: 80
+    # Data handling-related
+    batch_size_training: 50
+    batch_size_inference: 50
+
+predictor:
+  optim_kwargs:
+    loss_scaling: 1024
+
+trainer:
+  trainer:
+    accumulate_grad_batches: 4