Merge pull request #414 from WenkelF/explore_finetuning

Finetuning pipeline

README.md

@@ -85,21 +85,26 @@ If you are not familiar with [PyTorch](https://pytorch.org/docs) or [PyTorch-Lig
 ## 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
+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
 ```bash
-python main_run_multitask.py dataset=toymix model=gcn trainer.trainer.precision=32
+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
 ```bash
-python main_run_multitask.py dataset=toymix model=gcn accelerator=gpu
+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: 
+```bash
+graphium-train +finetuning=admet
+```
+
 To use a config file you built from scratch you can run
 ```bash
-python main_run_multitask.py --config-path [PATH] --config-name [CONFIG]
+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.
 

docs/cli_references.md

@@ -5,4 +5,5 @@ This page provides documentation for our command line tools.
 ::: mkdocs-click
     :module: graphium.cli
     :command: main_cli
-    :command: data_cli
+    :style: table
+    :prog_name: graphium

docs/tutorials/model_training/running-multitask-ipu.ipynb

@@ -420,7 +420,14 @@
     "logger.info(metrics)\n",
     "\n",
     "predictor = load_predictor(\n",
-    "    cfg, model_class, model_kwargs, metrics, accelerator_type, datamodule.task_norms\n",
+    "    cfg,\n",
+    "    model_class,\n",
+    "    model_kwargs,\n",
+    "    metrics,\n",
+    "    datamodule.get_task_levels(),\n",
+    "    accelerator_type,\n",
+    "    datamodule.featurization,\n",
+    "    datamodule.task_norms\n",
     ")\n",
     "logger.info(predictor.model)\n",
     "logger.info(ModelSummary(predictor, max_depth=4))"