Optimize framework performance with CSV buffering, configurable valid…

src/datasets/graph_dataset.py

@@ -11,10 +11,10 @@
 
 from datasets.utils.EdgeLabels import EdgeLabels
 from datasets.utils.NodeLabels import NodeLabels
-from src.datasets.graph_dataset_preprocessing import ZINCGraphDataPreprocessing, QMGraphDataPreprocessing, \
+from datasets.graph_dataset_preprocessing import ZINCGraphDataPreprocessing, QMGraphDataPreprocessing, \
     OGBGraphPropertyGraphDataPreprocessing, SubstructureBenchmarkPreprocessing, MergedGraphDataPreprocessing, \
     TUDatasetPreprocessing
-from src.utils.utils import load_graphs
+from utils.utils import load_graphs
 from torch_geometric.io import fs
 from torch_geometric.utils.convert import to_networkx
 from ogb.nodeproppred import PygNodePropPredDataset

src/framework/core.py

@@ -504,17 +504,12 @@ def collect_paths(main_configuration, model_configuration, dataset_configuration
 
 def copy_experiment_config(absolute_path, experiment_configuration, experiment_configuration_path,
                            graph_db_name):
+    import shutil
     results_path = experiment_configuration['paths']['results']
     if not results_path.joinpath(f"{graph_db_name}/config.yml"):
         source_path = Path(absolute_path).joinpath(experiment_configuration_path)
-        destination_path =results_path.joinpath(f"{graph_db_name}/config.yml")
-        # copy the config file to the results directory
-        # if linux
-        if os.name == 'posix':
-            os.system(f"cp {source_path} {destination_path}")
-        # if windows
-        elif os.name == 'nt':
-            os.system(f"copy {source_path} {destination_path}")
+        destination_path = results_path.joinpath(f"{graph_db_name}/config.yml")
+        shutil.copy2(source_path, destination_path)
 
 
 def preprocess_graph_data(experiment_configuration:dict):

src/framework/model_configuration.py

@@ -14,7 +14,7 @@
 from framework.utils.data_sampling import curriculum_sampling
 from framework.utils.parameters import Parameters
 from models import GraphModel
-from src.utils.utils import get_k_lowest_nonzero_indices, valid_pruning_configuration, is_pruning
+from utils.utils import get_k_lowest_nonzero_indices, valid_pruning_configuration, is_pruning
 from utils.timer import TimeClass
 
 
@@ -59,6 +59,8 @@ def __init__(self, run_id: int, k_val: int, graph_data: GraphDataset, model_data
         self.scheduler = None
         self.net = None
         self.class_weights = None
+        self._csv_buffer = []
+        self._csv_flush_interval = self.para.run_config.config.get('csv_flush_interval', 10)
         # get gpu or cpu: (cpu is recommended at the moment)
         if self.para.run_config.config.get('device', None) is not None:
             self.device = torch.device(self.para.run_config.config['device'] if torch.cuda.is_available() else "cpu")
@@ -102,6 +104,7 @@ def train_configuration(self, pretrained_network=None):
             # Test early stopping criterion
             if self.early_stopping(epoch):
                 print(f"Early stopping at epoch {epoch}")
+                self._flush_csv_buffer()
                 break
 
             timer.measure("epoch")
@@ -141,10 +144,12 @@ def train_configuration(self, pretrained_network=None):
                     self.scheduler.step()
 
             # Evaluate the results on training, validation and test set (only if specified in the config or for evaluation)
-            if self.validate_data.size != 0:
+            validation_frequency = self.para.run_config.config.get('validation_frequency', 1)
+            is_validation_epoch = (epoch + 1) % validation_frequency == 0 or epoch == self.para.n_epochs - 1
+            if is_validation_epoch and self.validate_data.size != 0:
                 epoch_values, validation_values, test_values = self.evaluate_results(epoch=epoch,train_values=epoch_values, validation_values=validation_values, test_values=test_values, evaluation_type='validation')
             # check wheter there is a test split
-            if self.test_data.size != 0:
+            if is_validation_epoch and self.test_data.size != 0:
                 epoch_values, validation_values, test_values = self.evaluate_results(epoch=epoch,train_values=epoch_values, validation_values=validation_values, test_values=test_values, evaluation_type='test')
 
             timer.measure("epoch")
@@ -775,11 +780,19 @@ def postprocess_writer(self, epoch, epoch_time, train_values: EvaluationValues,
                             f"{validation_values.accuracy};{validation_values.loss};" \
                             f"{test_values.accuracy};{test_values.loss}\n"
 
-        # Save file for results
+        # Buffer CSV writes and flush periodically
+        self._csv_buffer.append(res_str)
+        if len(self._csv_buffer) >= self._csv_flush_interval or epoch == self.para.n_epochs - 1:
+            self._flush_csv_buffer()
+
+    def _flush_csv_buffer(self):
+        if not self._csv_buffer:
+            return
         file_name = f'{self.para.db}_{self.para.config_id}_Results_run_id_{self.run_id}_validation_step_{self.para.validation_id}.csv'
         final_path = self.results_path.joinpath(f'{self.para.db}/Results/{file_name}')
         with open(final_path, "a") as file_obj:
-            file_obj.write(res_str)
+            file_obj.writelines(self._csv_buffer)
+        self._csv_buffer.clear()
 
 
     def train_graph_task(self, epoch, values, train_batches, timer):
@@ -856,9 +869,9 @@ def evaluate_graph_task(self, graph_ids):
 
         with torch.no_grad():
             self.net.train(False)
-            # Run ordinary GNN
-            # split the graph ids into batches of size 512 to avoid memory issues
-            batches = [graph_ids[i:i + 512] for i in range(0, len(graph_ids), 512)]
+            # split the graph ids into batches to avoid memory issues
+            eval_batch_size = self.para.run_config.config.get('eval_batch_size', 512)
+            batches = [graph_ids[i:i + eval_batch_size] for i in range(0, len(graph_ids), eval_batch_size)]
             # deduce output size from batches and graph data
             total_out_len = 0
             for batch in batches:
@@ -870,14 +883,14 @@ def evaluate_graph_task(self, graph_ids):
             loader = CustomBatchLoader(self.graph_data, batches)
             batch_counter = 0
             if not self.para.run_config.config.get('with_invariant_layers', True):
+                # Run ordinary GNN
                 for i, batch in enumerate(loader):
                     #print(f"Evaluating batch {i + 1}/{len(batches)}")
                     outputs[batch_counter:batch_counter + len(batch)] = self.net(batch_data=batch)
                     batch_counter += len(batch)
             else:
                 # Run Share GNN
                 for j, data_pos in enumerate(graph_ids):
-                    self.net.train(False)
                     outputs[j] = self.net(self.graph_data[data_pos], pos=data_pos)
         return labels, outputs
 

src/framework/utils/evaluation.py

@@ -4,7 +4,7 @@
 import pandas as pd
 from matplotlib import pyplot as plt
 
-from src.utils.utils import is_pruning
+from utils.utils import is_pruning
 
 
 def epoch_accuracy(db_name, y_val, ids):