Repository: datamol-io/graphium (language: python)

Issue:
Removed layers, simplified architectures
- Removal of unused GNNs
- Creation of ResidualConnection classes to avoid having multiple architectures. Then removed the architectures that relied on these residual connections
- Added some documentation and unit-tests
- Refactor PNA to have inheritance and avoid duplicated code
- Removed duplicated MLPs functions
- Changed the DGN networks to support both an input and output FeedForward network

This pull-request has a lot of progress in the refactoring of the code, but there is still a lot to do. 
The examples will no longer run after the merging, but I want to avoid pull-requests too large, so I think we can still go forward with the PR and do the remaining changes in another PR.


Discussion / hints:

Fix patch:
diff --git a/graphium/config/dummy_finetuning_from_gnn.yaml b/graphium/config/dummy_finetuning_from_gnn.yaml
new file mode 100644
index 000000000..3e7224b5e
--- /dev/null
+++ b/graphium/config/dummy_finetuning_from_gnn.yaml
@@ -0,0 +1,141 @@
+# Here, we are finetuning a FullGraphMultitaskNetwork
+# trained on ToyMix. We finetune from the gnn on the 
+# TDC dataset lipophilicity_astraceneca
+
+# Here are the changes to the architecture:
+#   Change gnn:
+#     depth:        4 -> 4 - 2 + 3 = 5
+#
+#   Keep modules after gnn and apply modifications
+#     graph_output_nn/graph:
+#       pooling: sum -> mean
+#       depth: 1 -> 2
+#     task_heads/zinc:
+#       new_sub_module: lipophilicity_astrazeneca
+#       out_dim: 3 -> 1
+#
+# Finetuning training:
+#   unfreeze one additional layer of pretrained gnn
+#   after 1 epochs, unfreeze all layers
+
+
+###################################################
+########### How to combine information  ###########
+###################################################
+
+
+###########################
+### FINETUNING-SPECIFIC ###
+###########################
+
+finetuning:
+  # New task
+  task: lipophilicity_astrazeneca
+  level: graph
+
+  # Pretrained model
+  pretrained_model_name: dummy-pretrained-model
+  finetuning_module: gnn 
+
+  # Changes to finetuning_module                                                  
+  drop_depth: 2
+  added_depth: 3
+
+  keep_modules_after_finetuning_module: # optional
+    graph_output_nn/graph:
+      pooling: [mean]
+      depth: 2
+    task_heads/zinc:
+      new_sub_module: lipophilicity_astrazeneca
+      out_dim: 1
+
+  # Finetuning training
+  unfreeze_pretrained_depth: 1
+  epoch_unfreeze_all: 1
+
+constants:
+  seed: 42
+  max_epochs: 2
+
+accelerator:
+  float32_matmul_precision: medium
+  type: cpu
+
+predictor:
+  random_seed: ${constants.seed}
+  optim_kwargs:
+    lr: 4.e-5
+  scheduler_kwargs: null
+  target_nan_mask: null
+  multitask_handling: flatten # flatten, mean-per-label
+  
+  torch_scheduler_kwargs:
+    module_type: WarmUpLinearLR
+    max_num_epochs: 2
+    warmup_epochs: 1
+    verbose: False
+  
+  metrics_on_progress_bar:
+    lipophilicity_astrazeneca: ["mae"]
+  loss_fun:
+    lipophilicity_astrazeneca: mae
+
+metrics:
+  lipophilicity_astrazeneca:
+    - name: mae
+      metric: mae
+      target_nan_mask: null
+      multitask_handling: flatten
+      threshold_kwargs: null
+    - name: spearman
+      metric: spearmanr
+      threshold_kwargs: null
+      target_nan_mask: null
+      multitask_handling: mean-per-label
+    - name: pearson
+      metric: pearsonr
+      threshold_kwargs: null
+      target_nan_mask: null
+      multitask_handling: mean-per-label
+    - name: r2_score
+      metric: r2
+      target_nan_mask: null
+      multitask_handling: mean-per-label
+      threshold_kwargs: null
+
+trainer:
+  seed: ${constants.seed}
+  trainer:
+    precision: 32
+    max_epochs: 2
+    min_epochs: 1
+    check_val_every_n_epoch: 1
+    accumulate_grad_batches: 1
+  
+##################
+### DATAMODULE ###
+##################
+
+datamodule:
+
+### FROM FINETUNING ###
+
+  module_type: "ADMETBenchmarkDataModule"
+  args:
+    # TDC specific
+    tdc_benchmark_names: [lipophilicity_astrazeneca]
+    tdc_train_val_seed: ${constants.seed}
+    
+    batch_size_training: 200
+    batch_size_inference: 200
+    featurization_n_jobs: 0
+    num_workers: 0
+
+    prepare_dict_or_graph: pyg:graph
+    featurization_progress: True
+    featurization_backend: "loky"
+    persistent_workers: False
+
+
+
+
diff --git a/graphium/config/dummy_finetuning.yaml b/graphium/config/dummy_finetuning_from_task_head.yaml
similarity index 89%
rename from graphium/config/dummy_finetuning.yaml
rename to graphium/config/dummy_finetuning_from_task_head.yaml
index 38cdf029b..27b56a683 100644
--- a/graphium/config/dummy_finetuning.yaml
+++ b/graphium/config/dummy_finetuning_from_task_head.yaml
@@ -3,16 +3,18 @@
 # (graph-level) on the TDC dataset lipophilicity_astraceneca
 
 # Here are the changes to the architecture:
+#   Change zinc task-head:
+#     depth:        2 -> 2 - 1 + 2 = 3
+#     out_dim:      3 -> 8
 #
-# Change zinc task-head:
-#   depth:        2 -> 2 - 1 + 2 = 3
-#   out_dim:      3 -> 8
+#   Add finetuning head
+#     model_type:   FeedForwardNN
+#     out_dim:      1
+#     hidden_dims:  8
+#     depth:        2
 #
-# Add finetuning head
-#   model_type:   FeedForwardNN
-#   out_dim:      1
-#   hidden_dims:  8
-#   depth:        2
+# Finetuning training:
+#   after 1 epochs, unfreeze all layers
 
 
 ###################################################
@@ -55,11 +57,11 @@ finetuning:
 
   # Finetuning training
   unfreeze_pretrained_depth: 0
-  epoch_unfreeze_all: 2
+  epoch_unfreeze_all: 1
 
 constants:
   seed: 42
-  max_epochs: 3
+  max_epochs: 2
 
 accelerator:
   float32_matmul_precision: medium
@@ -75,7 +77,7 @@ predictor:
   
   torch_scheduler_kwargs:
     module_type: WarmUpLinearLR
-    max_num_epochs: 3
+    max_num_epochs: 2
     warmup_epochs: 1
     verbose: False
   
@@ -111,7 +113,7 @@ trainer:
   seed: ${constants.seed}
   trainer:
     precision: 32
-    max_epochs: 3
+    max_epochs: 2
     min_epochs: 1
     check_val_every_n_epoch: 1
     accumulate_grad_batches: 1
@@ -138,7 +140,6 @@ datamodule:
     prepare_dict_or_graph: pyg:graph
     featurization_progress: True
     featurization_backend: "loky"
-    processed_graph_data_path: "../datacache/neurips2023-small/"
     persistent_workers: False
 
 


Test patch:
diff --git a/tests/test_finetuning.py b/tests/test_finetuning.py
index 7eeaf5f3f..95ed3ddff 100644
--- a/tests/test_finetuning.py
+++ b/tests/test_finetuning.py
@@ -34,7 +34,7 @@
 
 
 class Test_Finetuning(ut.TestCase):
-    def test_finetuning_pipeline(self):
+    def test_finetuning_from_task_head(self):
         # Skip test if PyTDC package not installed
         try:
             import tdc
@@ -45,7 +45,7 @@ def test_finetuning_pipeline(self):
         ### Test modification of config for finetuning ###
         ##################################################
 
-        cfg = graphium.load_config(name="dummy_finetuning")
+        cfg = graphium.load_config(name="dummy_finetuning_from_task_head")
         cfg = OmegaConf.to_container(cfg, resolve=True)
 
         cfg = modify_cfg_for_finetuning(cfg)
@@ -55,6 +55,7 @@ def test_finetuning_pipeline(self):
 
         # Load and initialize the dataset
         datamodule = load_datamodule(cfg, accelerator_type)
+        datamodule.task_specific_args["lipophilicity_astrazeneca"].sample_size = 100
 
         # Initialize the network
         model_class, model_kwargs = load_architecture(
@@ -77,15 +78,21 @@ def test_finetuning_pipeline(self):
             datamodule.task_norms,
         )
 
-        self.assertEqual(
-            len(
-                predictor.model.pretrained_model.net.task_heads.task_heads["lipophilicity_astrazeneca"].layers
-            ),
-            3,
+        # Create module map
+        module_map = deepcopy(predictor.model.pretrained_model.net._module_map)
+
+        cfg_finetune = cfg["finetuning"]
+        finetuning_module = "".join([cfg_finetune["finetuning_module"], "/", cfg_finetune["task"]])
+        finetuning_module_from_pretrained = "".join(
+            [cfg_finetune["finetuning_module"], "/", cfg_finetune["sub_module_from_pretrained"]]
         )
+
+        # Test for correctly modified shapes and number of layers in finetuning module
         self.assertEqual(
-            predictor.model.pretrained_model.net.task_heads.task_heads["lipophilicity_astrazeneca"].out_dim, 8
+            len(module_map[finetuning_module]),
+            3,
         )
+        self.assertEqual(module_map[finetuning_module][-1].linear.weight.size(0), 8)
         self.assertEqual(predictor.model.finetuning_head.net.in_dim, 8)
         self.assertEqual(len(predictor.model.finetuning_head.net.layers), 2)
         self.assertEqual(predictor.model.finetuning_head.net.out_dim, 1)
@@ -101,40 +108,34 @@ def test_finetuning_pipeline(self):
 
         pretrained_model.create_module_map()
         module_map_from_pretrained = deepcopy(pretrained_model._module_map)
-        module_map = deepcopy(predictor.model.pretrained_model.net._module_map)
 
         # Finetuning module has only been partially overwritten
-        cfg_finetune = cfg["finetuning"]
-        finetuning_module = "".join([cfg_finetune["finetuning_module"], "/", cfg_finetune["task"]])
-        finetuning_module_from_pretrained = "".join(
-            [cfg_finetune["finetuning_module"], "/", cfg_finetune["sub_module_from_pretrained"]]
-        )
+        loaded_layers = module_map_from_pretrained[finetuning_module_from_pretrained]
+        overwritten_layers = module_map[finetuning_module]
 
-        pretrained_layers = module_map[finetuning_module]
-        overwritten_layers = module_map_from_pretrained[finetuning_module_from_pretrained]
-
-        for idx, (pretrained, overwritten) in enumerate(zip(pretrained_layers, overwritten_layers)):
+        for idx, (loaded, overwritten) in enumerate(zip(loaded_layers, overwritten_layers)):
             if idx < 1:
-                assert torch.equal(pretrained.linear.weight, overwritten.linear.weight)
-                assert torch.equal(pretrained.linear.bias, overwritten.linear.bias)
+                assert torch.equal(loaded.linear.weight, overwritten.linear.weight)
+                assert torch.equal(loaded.linear.bias, overwritten.linear.bias)
             else:
-                assert not torch.equal(pretrained.linear.weight, overwritten.linear.weight)
-                assert not torch.equal(pretrained.linear.bias, overwritten.linear.bias)
+                assert not torch.equal(loaded.linear.weight, overwritten.linear.weight)
+                assert not torch.equal(loaded.linear.bias, overwritten.linear.bias)
 
-            if idx + 1 == min(len(pretrained_layers), len(overwritten_layers)):
+            if idx + 1 == min(len(loaded_layers), len(overwritten_layers)):
                 break
 
-        _ = module_map.popitem(last=True)
-        overwritten_modules = module_map.values()
-
-        _ = module_map_from_pretrained.popitem(last=True)
-        pretrained_modules = module_map_from_pretrained.values()
+        for module_name in module_map.keys():
+            if module_name == finetuning_module:
+                break
 
-        for overwritten_module, pretrained_module in zip(overwritten_modules, pretrained_modules):
-            for overwritten, pretrained in zip(
-                overwritten_module.parameters(), pretrained_module.parameters()
+            loaded_module, overwritten_module = (
+                module_map_from_pretrained[module_name],
+                module_map[module_name],
+            )
+            for loaded_params, overwritten_params in zip(
+                loaded_module.parameters(), overwritten_module.parameters()
             ):
-                assert torch.equal(overwritten.data, pretrained.data)
+                assert torch.equal(loaded_params.data, overwritten_params.data)
 
         #################################################
         ### Test correct (un)freezing during training ###
@@ -222,6 +223,181 @@ def on_train_epoch_start(self, trainer, pl_module):
         # Run the model training
         trainer.fit(model=predictor, datamodule=datamodule)
 
+    def test_finetuning_from_gnn(self):
+        # Skip test if PyTDC package not installed
+        try:
+            import tdc
+        except ImportError:
+            self.skipTest("PyTDC needs to be installed to run this test. Use `pip install PyTDC`.")
+
+        ##################################################
+        ### Test modification of config for finetuning ###
+        ##################################################
+
+        cfg = graphium.load_config(name="dummy_finetuning_from_gnn")
+        cfg = OmegaConf.to_container(cfg, resolve=True)
+
+        cfg = modify_cfg_for_finetuning(cfg)
+
+        # Initialize the accelerator
+        cfg, accelerator_type = load_accelerator(cfg)
+
+        # Load and initialize the dataset
+        datamodule = load_datamodule(cfg, accelerator_type)
+        datamodule.task_specific_args["lipophilicity_astrazeneca"].sample_size = 100
+
+        # Initialize the network
+        model_class, model_kwargs = load_architecture(
+            cfg,
+            in_dims=datamodule.in_dims,
+        )
+
+        datamodule.prepare_data()
+
+        metrics = load_metrics(cfg)
+
+        predictor = load_predictor(
+            cfg,
+            model_class,
+            model_kwargs,
+            metrics,
+            datamodule.get_task_levels(),
+            accelerator_type,
+            datamodule.featurization,
+            datamodule.task_norms,
+        )
+
+        # Create module map
+        module_map = deepcopy(predictor.model.pretrained_model.net._module_map)
+
+        cfg_finetune = cfg["finetuning"]
+        finetuning_module = cfg_finetune["finetuning_module"]
+
+        # Test for correctly modified shapes and number of layers in finetuning module
+        self.assertEqual(
+            len(module_map[finetuning_module]),
+            5,
+        )
+        self.assertEqual(module_map[finetuning_module][-1].model.lin.weight.size(0), 96)
+        self.assertEqual(len(module_map["graph_output_nn/graph"]), 2)
+
+        assert predictor.model.pretrained_model.net.task_heads.graph_output_nn[
+            "graph"
+        ].graph_output_nn_kwargs["graph"]["pooling"] == ["mean"]
+
+        ################################################
+        ### Test overwriting with pretrained weights ###
+        ################################################
+
+        # Load pretrained & replace in predictor
+        pretrained_model = PredictorModule.load_pretrained_models(
+            cfg["finetuning"]["pretrained_model_name"], device="cpu"
+        ).model
+
+        pretrained_model.create_module_map()
+        module_map_from_pretrained = deepcopy(pretrained_model._module_map)
+
+        # Finetuning module has only been partially overwritten
+        loaded_layers = module_map_from_pretrained[finetuning_module]
+        overwritten_layers = module_map[finetuning_module]
+
+        for idx, (loaded, overwritten) in enumerate(zip(loaded_layers, overwritten_layers)):
+            if idx < 2:
+                assert torch.equal(loaded.model.lin.weight, overwritten.model.lin.weight)
+            else:
+                assert not torch.equal(loaded.model.lin.weight, overwritten.model.lin.weight)
+
+            if idx + 1 == min(len(loaded_layers), len(overwritten_layers)):
+                break
+
+        for module_name in module_map.keys():
+            if module_name == finetuning_module:
+                break
+
+            loaded_module, overwritten_module = (
+                module_map_from_pretrained[module_name],
+                module_map[module_name],
+            )
+            for loaded_params, overwritten_params in zip(
+                loaded_module.parameters(), overwritten_module.parameters()
+            ):
+                assert torch.equal(loaded_params.data, overwritten_params.data)
+
+        #################################################
+        ### Test correct (un)freezing during training ###
+        #################################################
+
+        # Define test callback that checks for correct (un)freezing
+        class TestCallback(Callback):
+            def __init__(self, cfg):
+                super().__init__()
+
+                self.cfg_finetune = cfg["finetuning"]
+
+            def on_train_epoch_start(self, trainer, pl_module):
+                module_map = pl_module.model.pretrained_model.net._module_map
+
+                training_depth = self.cfg_finetune["added_depth"] + self.cfg_finetune.pop(
+                    "unfreeze_pretrained_depth", 0
+                )
+
+                frozen_parameters, unfrozen_parameters = [], []
+
+                if trainer.current_epoch == 0:
+                    frozen = True
+
+                    for module_name, module in module_map.items():
+                        if module_name == finetuning_module:
+                            # After the finetuning module, all parameters are unfrozen
+                            frozen = False
+
+                            frozen_parameters.extend(
+                                [
+                                    parameter.requires_grad
+                                    for parameter in module[:-training_depth].parameters()
+                                ]
+                            )
+                            unfrozen_parameters.extend(
+                                [
+                                    parameter.requires_grad
+                                    for parameter in module[-training_depth:].parameters()
+                                ]
+                            )
+                            continue
+
+                        if frozen:
+                            frozen_parameters.extend(
+                                [parameter.requires_grad for parameter in module.parameters()]
+                            )
+                        else:
+                            unfrozen_parameters.extend(
+                                [parameter.requires_grad for parameter in module.parameters()]
+                            )
+
+                    assert not True in frozen_parameters
+                    assert not False in unfrozen_parameters
+
+                if trainer.current_epoch == 1:
+                    # All parameter are unfrozen starting from epoch_unfreeze_all
+                    unfrozen_parameters = [
+                        parameter.requires_grad for parameter in pl_module.model.parameters()
+                    ]
+
+                    assert not False in unfrozen_parameters
+
+        trainer = load_trainer(cfg, accelerator_type)
+
+        finetuning_training_kwargs = cfg["finetuning"]["training_kwargs"]
+        trainer.callbacks.append(GraphFinetuning(**finetuning_training_kwargs))
+
+        # Add test callback to trainer
+        trainer.callbacks.append(TestCallback(cfg))
+
+        predictor.set_max_nodes_edges_per_graph(datamodule, stages=["train", "val"])
+
+        # Run the model training
+        trainer.fit(model=predictor, datamodule=datamodule)
+
 
 if __name__ == "__main__":
     ut.main()