some config changes

expts/neurips2023_configs/config_large_gcn.yaml

@@ -25,7 +25,7 @@ accelerator:
         loss_scaling: 1024
     trainer:
       trainer:
-        precision: 16
+        precision: 32
         accumulate_grad_batches: 8
 
   ipu_config:
@@ -97,6 +97,7 @@ datamodule:
         task_level: graph
         splits_path: graphium/data/neurips2023/large-dataset/pcqm4m_g25_n4_random_splits.pt  # Download with `wget https://storage.googleapis.com/graphium-public/datasets/neurips_2023/Large-dataset/pcqm4m_g25_n4_random_splits.pt`
         label_normalization:
+          normalize_val_test: True
           method: "normal"
 
       pcqm4m_n4:
@@ -111,14 +112,15 @@ datamodule:
         splits_path: graphium/data/neurips2023/large-dataset/pcqm4m_g25_n4_random_splits.pt  # Download with `wget https://storage.googleapis.com/graphium-public/datasets/neurips_2023/Large-dataset/pcqm4m_g25_n4_random_splits.pt`
         seed: *seed
         label_normalization:
+          normalize_val_test: True
           method: "normal"
 
     # Featurization
     prepare_dict_or_graph: pyg:graph
     featurization_n_jobs: 30
     featurization_progress: True
     featurization_backend: "loky"
-    processed_graph_data_path: "../neurips2023-large/"
+    processed_graph_data_path: "/net/group/software-apps/datacache/neurips2023-large/"
     featurization:
     # OGB: ['atomic_num', 'degree', 'possible_formal_charge', 'possible_numH' (total-valence),
     # 'possible_number_radical_e', 'possible_is_aromatic', 'possible_is_in_ring',
@@ -379,6 +381,11 @@ metrics:
       threshold_kwargs: null
       target_nan_mask: null
       multitask_handling: mean-per-label
+    - name: r2
+      metric: r2_score_ipu
+      threshold_kwargs: null
+      target_nan_mask: null
+      multitask_handling: mean-per-label
   pcqm4m_n4: *pcqm_metrics
 
 trainer:

expts/neurips2023_configs/config_large_gin.yaml

@@ -11,11 +11,11 @@ accelerator:
       args:
         ipu_dataloader_training_opts:
           mode: async
-          max_num_nodes_per_graph: 45 # train max nodes: 20, max_edges: 54
-          max_num_edges_per_graph: 95
+          max_num_nodes_per_graph: 60 # train max nodes: 20, max_edges: 54
+          max_num_edges_per_graph: 100
         ipu_dataloader_inference_opts:
           mode: async
-          max_num_nodes_per_graph: 50 # valid max nodes: 51, max_edges: 118
+          max_num_nodes_per_graph: 60 # valid max nodes: 51, max_edges: 118
           max_num_edges_per_graph: 100
         # Data handling-related
         batch_size_training: 10
@@ -97,6 +97,7 @@ datamodule:
         task_level: graph
         splits_path: graphium/data/neurips2023/large-dataset/pcqm4m_g25_n4_random_splits.pt  # Download with `wget https://storage.googleapis.com/graphium-public/datasets/neurips_2023/Large-dataset/pcqm4m_g25_n4_random_splits.pt`
         label_normalization:
+          normalize_val_test: True
           method: "normal"
 
       pcqm4m_n4:
@@ -111,14 +112,15 @@ datamodule:
         splits_path: graphium/data/neurips2023/large-dataset/pcqm4m_g25_n4_random_splits.pt  # Download with `wget https://storage.googleapis.com/graphium-public/datasets/neurips_2023/Large-dataset/pcqm4m_g25_n4_random_splits.pt`
         seed: *seed
         label_normalization:
+          normalize_val_test: True
           method: "normal"
 
     # Featurization
     prepare_dict_or_graph: pyg:graph
     featurization_n_jobs: 30
     featurization_progress: True
     featurization_backend: "loky"
-    processed_graph_data_path: "../neurips2023-large/"
+    processed_graph_data_path: "/net/group/software-apps/datacache/neurips2023-large/"
     featurization:
     # OGB: ['atomic_num', 'degree', 'possible_formal_charge', 'possible_numH' (total-valence),
     # 'possible_number_radical_e', 'possible_is_aromatic', 'possible_is_in_ring',
@@ -373,6 +375,11 @@ metrics:
       threshold_kwargs: null
       target_nan_mask: null
       multitask_handling: mean-per-label
+    - name: r2
+      metric: r2_score_ipu
+      threshold_kwargs: null
+      target_nan_mask: null
+      multitask_handling: mean-per-label
   pcqm4m_n4: *pcqm_metrics
 
 trainer:

expts/neurips2023_configs/single_task/config_large_gine_pcba.yaml

@@ -15,8 +15,8 @@ accelerator:
           max_num_edges_per_graph: 100
         ipu_dataloader_inference_opts:
           mode: async
-          max_num_nodes_per_graph: 60 # valid max nodes: 51, max_edges: 118
-          max_num_edges_per_graph: 100
+          max_num_nodes_per_graph: 90 # valid max nodes: 51, max_edges: 118
+          max_num_edges_per_graph: 300
         # Data handling-related
         batch_size_training: 10
         batch_size_inference: 2

expts/neurips2023_configs/single_task/config_large_gine_pcq.yaml

@@ -19,7 +19,7 @@ accelerator:
           max_num_edges_per_graph: 100
         # Data handling-related
         batch_size_training: 10
-        batch_size_inference: 2
+        batch_size_inference: 10
     predictor:
       optim_kwargs:
         loss_scaling: 1024

graphium/trainer/predictor.py

@@ -347,8 +347,8 @@ def _general_step(self, batch: Dict[str, Tensor], step_name: str, to_cpu: bool)
                 # if normalize_val_test is true, no denormalization is applied, all losses and metrics are normalized version
                 preds[task] = task_specific_norm.denormalize(preds[task])
                 targets_dict[task] = task_specific_norm.denormalize(targets_dict[task])
-            preds[task] = preds[task].clone().detach().to(device=device)
-            targets_dict[task] = targets_dict[task].clone().detach().to(device=device)
+            preds[task] = preds[task].detach().to(device=device)
+            targets_dict[task] = targets_dict[task].detach().to(device=device)
         if weights is not None:
             weights = weights.detach().to(device=device)
 
@@ -447,17 +447,17 @@ def on_train_batch_end(self, outputs, batch: Any, batch_idx: int) -> None:
         tput = num_graphs / train_batch_time
 
         # this code is likely repeated for validation and testing, this should be moved to a function
-        # self.task_epoch_summary.update_predictor_state(
-        #     step_name="train",
-        #     targets=outputs["targets"],
-        #     predictions=outputs["preds"],
-        #     loss=outputs["loss"],  # This is the weighted loss for now, but change to task-specific loss
-        #     task_losses=outputs["task_losses"],
-        #     n_epochs=self.current_epoch,
-        # )
-        # metrics_logs = self.task_epoch_summary.get_metrics_logs()  # Dict[task, metric_logs]
-        # metrics_logs["_global"]["grad_norm"] = self.get_gradient_norm()
-        # outputs.update(metrics_logs)  # Dict[task, metric_logs]. Concatenate them?
+        self.task_epoch_summary.update_predictor_state(
+            step_name="train",
+            targets=outputs["targets"],
+            predictions=outputs["preds"],
+            loss=outputs["loss"],  # This is the weighted loss for now, but change to task-specific loss
+            task_losses=outputs["task_losses"],
+            n_epochs=self.current_epoch,
+        )
+        metrics_logs = self.task_epoch_summary.get_metrics_logs()  # Dict[task, metric_logs]
+        metrics_logs["_global"]["grad_norm"] = self.get_gradient_norm()
+        outputs.update(metrics_logs)  # Dict[task, metric_logs]. Concatenate them?
 
         concatenated_metrics_logs = {}  # self.task_epoch_summary.concatenate_metrics_logs(metrics_logs)
         concatenated_metrics_logs["loss"] = outputs["loss"]
@@ -487,8 +487,8 @@ def training_step(self, batch: Dict[str, Tensor], to_cpu: bool = True) -> Dict[s
             # step_dict = self._general_step(batch=batch, step_name="train", to_cpu=True)
             step_dict = self._general_step(batch=batch, step_name="train", to_cpu=to_cpu)
 
-        step_dict.pop("preds")
-        step_dict.pop("targets")
+        # step_dict.pop("preds")
+        # step_dict.pop("targets")
 
         return step_dict  # Returning the metrics_logs with the loss