From cbc590a48c2577ba051afaaf852133ec6bc7d725 Mon Sep 17 00:00:00 2001 From: florian Date: Sat, 7 Feb 2026 09:21:10 +0100 Subject: [PATCH 1/4] Optimize framework performance with CSV buffering, configurable validation frequency, and improved file operations - Replace OS-specific commands with shutil.copy2 for cross-platform file copying - Add CSV buffering mechanism to reduce I/O operations (flush every 50 writes) - Add configurable validation frequency to skip validation on some epochs - Optimize eval/train mode switching to avoid redundant calls - Make evaluation batch size configurable (default 512) - Remove redundant self.net.train(False) call in evaluation loop Co-Authored-By: Claude Sonnet 4.5 --- src/framework/core.py | 11 +++------ src/framework/model_configuration.py | 34 ++++++++++++++++++++-------- 2 files changed, 28 insertions(+), 17 deletions(-) diff --git a/src/framework/core.py b/src/framework/core.py index 188a43e..93a075a 100644 --- a/src/framework/core.py +++ b/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): diff --git a/src/framework/model_configuration.py b/src/framework/model_configuration.py index bf3adb0..e3db69e 100644 --- a/src/framework/model_configuration.py +++ b/src/framework/model_configuration.py @@ -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 = 50 # 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") @@ -348,7 +353,8 @@ def evaluate_results(self, epoch: int, batch_length=0, num_batches=0, batches=None): - self.net.eval() + if evaluation_type != 'training': + self.net.eval() if evaluation_type == 'training': batch_acc = 0 # if num classes is one calculate the mae and mae_std or if the task is regression @@ -633,7 +639,8 @@ def evaluate_results(self, epoch: int, labels_np = labels_np.T df = pd.DataFrame(labels_np) df.to_csv("Results/Parameter/test_predictions.csv", header=False, index=False, mode='a') - self.net.train() + if evaluation_type != 'training': + self.net.train() return train_values, validation_values, test_values def preprocess_writer(self)-> bool: @@ -775,11 +782,20 @@ 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): @@ -857,8 +873,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: @@ -877,7 +894,6 @@ def evaluate_graph_task(self, graph_ids): 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 From e1938387ffa45a8bf0e8ed1557774bd626a50cd0 Mon Sep 17 00:00:00 2001 From: florian Date: Sat, 7 Feb 2026 09:51:53 +0100 Subject: [PATCH 2/4] Add additional model layer optimizations and security fixes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit **Performance optimizations:** - Replace torch.nn.Dropout() with torch.nn.functional.dropout() in all MPNN layers for better performance - Optimize ShareGNN weight distribution using chunks instead of preallocated tensors - Remove redundant code and unused variables (representation_list, n_heads_per_label, num field) - Simplify num_heads calculation in framework_layer **Security fixes:** - Replace unsafe eval() with ACTIVATION_MAP dictionary for activation functions - Replace eval() with ast.literal_eval() in ShareGNN preprocessing **Code quality improvements:** - Fix inconsistent return patterns in model layer creation - Remove unused pandas import in inv_based_pooling - Fix property_dict bug in ShareGNN utils - Remove redundant out_features initialization in share_gnn_linear - Improve variable naming consistency in layer_loader (head → heads) Co-Authored-By: Claude Sonnet 4.5 --- .../layers/inv_based_message_passing.py | 101 +++++++++--------- .../ShareGNN/layers/inv_based_pooling.py | 53 +++++---- .../ShareGNN/layers/share_gnn_linear.py | 3 - .../ShareGNN/preprocessing/preprocessing.py | 3 +- src/models/ShareGNN/utils.py | 3 +- src/models/layers/framework_layer.py | 27 ++++- src/models/layers/mpnn_classical/gat_conv.py | 2 +- .../layers/mpnn_classical/gatv2_conv.py | 2 +- src/models/layers/mpnn_classical/gcn_conv.py | 2 +- src/models/layers/mpnn_classical/gin_conv.py | 2 +- src/models/layers/mpnn_classical/sage_conv.py | 2 +- src/models/layers/utils/layer_loader.py | 32 +++--- src/models/model.py | 16 +-- 13 files changed, 135 insertions(+), 113 deletions(-) diff --git a/src/models/ShareGNN/layers/inv_based_message_passing.py b/src/models/ShareGNN/layers/inv_based_message_passing.py index 676bae7..a836f75 100644 --- a/src/models/ShareGNN/layers/inv_based_message_passing.py +++ b/src/models/ShareGNN/layers/inv_based_message_passing.py @@ -47,7 +47,6 @@ def __init__(self, parameters: Parameters, layer: Layer, graph_data: GraphDatase super(InvariantBasedMessagePassingLayer, self).__init__(parameters, layer, graph_data) self.out_features = self.in_features * self.num_heads - self.n_heads_per_label = [] # number of heads per node label description for h_id, head in enumerate(layer.layer_heads): self.source_label_descriptions.append(layer.get_source_string(h_id)) @@ -58,35 +57,30 @@ def __init__(self, parameters: Parameters, layer: Layer, graph_data: GraphDatase self.n_bias_labels.append(graph_data.node_labels[self.bias_label_descriptions[h_id]].num_unique_node_labels) self.property_descriptions.append(head.property_dict.get_property_string()) self.n_properties.append(graph_data.properties[self.property_descriptions[h_id]].num_properties[(self.layer_id, h_id)]) - self.n_heads_per_label.append(head.num) self.bias_list = [head.bias for head in layer.layer_heads] self.bias = any(self.bias_list) # check if bias is used - - # Determine the number of weights and biases # There are two cases asymetric and symmetric, asymetric is the default, TODO add symmetric case self.skips = [0] self.skips_description = [None] self.skips_description_text = [None] - weight_dist_default_size = 1000000 - self.weight_distribution = torch.zeros((weight_dist_default_size, 4), dtype=torch.int64).to(self.device) - last_weight_distribution_index = 0 - self.bias_distribution = [None] * len(graph_data) + weight_distribution_chunks = [[] for _ in range(len(graph_data))] + bias_distribution_chunks = [[] for _ in range(len(graph_data))] # Iterate over all heads in the layer - for head_id, head in enumerate(self.layer.layer_heads): + for i, head in enumerate(self.layer.layer_heads): # get all the valid property values for the head (e.g., the distances 0, 3, 6) - valid_property_values = self.graph_data.properties[self.property_descriptions[head_id]].valid_values[(self.layer_id, head_id)] + valid_property_values = self.graph_data.properties[self.property_descriptions[i]].valid_values[(self.layer_id, i)] # apply the head and tail labels to the subdict - source_labels = self.graph_data.node_labels[self.source_label_descriptions[head_id]].node_labels - target_labels = self.graph_data.node_labels[self.target_label_descriptions[head_id]].node_labels - bias_labels = self.graph_data.node_labels[self.bias_label_descriptions[head_id]].node_labels + source_labels = self.graph_data.node_labels[self.source_label_descriptions[i]].node_labels + target_labels = self.graph_data.node_labels[self.target_label_descriptions[i]].node_labels + bias_labels = self.graph_data.node_labels[self.bias_label_descriptions[i]].node_labels for key in valid_property_values: #print(f'Initialize head {i+1}/{len(self.layer.layer_heads)} with property {key}') - property_subdict = self.graph_data.properties[self.property_descriptions[head_id]].properties[key] - property_subdict_slices = self.graph_data.properties[self.property_descriptions[head_id]].properties_slices[key] + property_subdict = self.graph_data.properties[self.property_descriptions[i]].properties[key] + property_subdict_slices = self.graph_data.properties[self.property_descriptions[i]].properties_slices[key] labeled_subdict = property_subdict.detach().clone() labeled_subdict[:, 0] = source_labels[property_subdict[:, 0]] labeled_subdict[:, 1] = target_labels[property_subdict[:, 1]] @@ -120,26 +114,27 @@ def __init__(self, parameters: Parameters, layer: Layer, graph_data: GraphDatase indices[valid_indices] = torch.tensor([valid_value_dict[idx.item()] for idx in indices[valid_indices]], dtype=torch.int64) num_weights = len(valid_values) start_time = time.time() - for j in range(head.num): - for idx in range(len(graph_data)): - # if number of graphs is larger than 10000 print progress - if len(graph_data) > 10000 and idx % 1000 == 0: - print(f'Head {head_id+1}/{len(self.layer.layer_heads)} with property {key}: {idx}/{len(graph_data)} graphs processed ({(idx/len(graph_data))*100:.2f}%) time so far (in s): {time.time()-start_time:.2f}', flush=True) - # get the valid indices for the current graph - if threshold > 1 or do_invalid_indices_exist or upper_threshold is not None: - valid_indices_graph = torch.where(valid_indices_bool[property_subdict_slices[idx]:property_subdict_slices[idx+1]])[0] + property_subdict_slices[idx] - else: - valid_indices_graph = torch.arange(property_subdict_slices[idx], property_subdict_slices[idx+1], dtype=torch.int64) - # create new tensor where each row is the concatenation of head_id, property_subdict_row, and indices - - # Use self.weight_distribution directly - self.weight_distribution[last_weight_distribution_index:last_weight_distribution_index + len(valid_indices_graph), 0] = head_id - self.weight_distribution[last_weight_distribution_index:last_weight_distribution_index + len(valid_indices_graph), 1:3] = property_subdict[valid_indices_graph] - self.graph_data.slices['x'][idx] # check if subtracting is necessary - self.weight_distribution[last_weight_distribution_index:last_weight_distribution_index + len(valid_indices_graph), 3] = indices[valid_indices_graph] + self.skips[-1] + for idx in range(len(graph_data)): + # if number of graphs is larger than 10000 print progress + if len(graph_data) > 10000 and idx % 1000 == 0: + print(f'Head {i+1}/{len(self.layer.layer_heads)} with property {key}: {idx}/{len(graph_data)} graphs processed ({(idx/len(graph_data))*100:.2f}%) time so far (in s): {time.time()-start_time:.2f}', flush=True) + # get the valid indices for the current graph + if threshold > 1 or do_invalid_indices_exist or upper_threshold is not None: + valid_indices_graph = torch.where(valid_indices_bool[property_subdict_slices[idx]:property_subdict_slices[idx+1]])[0] + property_subdict_slices[idx] + else: + valid_indices_graph = torch.arange(property_subdict_slices[idx], property_subdict_slices[idx+1], dtype=torch.int64) + # create new tensor where each row is the concatenation of head_id, property_subdict_row, and indices + new_weight_distribution = torch.zeros((len(valid_indices_graph), 4), dtype=torch.int64) + new_weight_distribution[:, 0] = i + new_weight_distribution[:, 1:3] = property_subdict[valid_indices_graph] - self.graph_data.slices['x'][idx] # check if subtracting is necessary + new_weight_distribution[:, 3] = indices[valid_indices_graph] + self.skips[-1] + weight_distribution_chunks[idx].append(new_weight_distribution) + + self.skips.append(self.skips[-1] + num_weights) - self.skips_description.append({'head:': head_id, 'property': key, 'weights': num_weights}) - self.skips_description_text.append(f"Head {head_id} Property {key} has {num_weights} different weights") + self.skips_description.append({'head:': i, 'property': key, 'weights': num_weights}) + self.skips_description_text.append(f"Head {i} Property {key} has {num_weights} different weights") self.weight_num.append(self.skips[-1]) @@ -147,28 +142,34 @@ def __init__(self, parameters: Parameters, layer: Layer, graph_data: GraphDatase if self.bias: # Determine the number of different learnable parameters in the bias vector - self.bias_num.append(self.in_features * self.n_bias_labels[head_id]) + self.bias_num.append(self.in_features * self.n_bias_labels[i]) # Set the bias weights _, indices, counts = torch.unique(bias_labels, dim=0, return_inverse=True, return_counts=True, sorted=False) for idx in range(len(graph_data)): for feature_id in range(self.in_features): new_bias_distribution = torch.zeros((graph_data.num_nodes[idx].item(), 4), dtype=torch.int64) - new_bias_distribution[:, 0] = head_id + new_bias_distribution[:, 0] = i new_bias_distribution[:, 1] = torch.arange(graph_data.num_nodes[idx].item(), dtype=torch.int64) # alternative torch.arange(start=graph_data.slices['x'][idx], end=graph_data.slices['x'][idx+1], dtype=torch.int64) new_bias_distribution[:, 2] = feature_id - new_bias_distribution[:, 3] = indices[graph_data.slices['x'][idx]:graph_data.slices['x'][idx+1]] + feature_id * self.n_bias_labels[head_id] - if self.bias_distribution[idx] is None: - self.bias_distribution[idx] = new_bias_distribution.detach().clone() - else: - self.bias_distribution[idx] = torch.cat((self.bias_distribution[idx], new_bias_distribution), dim=0) + new_bias_distribution[:, 3] = indices[graph_data.slices['x'][idx]:graph_data.slices['x'][idx+1]] + feature_id * self.n_bias_labels[i] + bias_distribution_chunks[idx].append(new_bias_distribution) # Merge the weight distribution of all graphs (creating additionally slicing information) - self.weight_distribution_slices = torch.tensor([0] + [len(w) for w in self.weight_distribution], dtype=torch.int64).cumsum(dim=0) - self.weight_distribution = torch.cat([self.weight_distribution[i] for i in range(len(graph_data))], dim=0).to(self.device) + # Single torch.cat per graph instead of repeated incremental concatenation + merged_weight_distributions = [ + torch.cat(chunks, dim=0) if chunks else torch.zeros((0, 4), dtype=torch.int64) + for chunks in weight_distribution_chunks + ] + self.weight_distribution_slices = torch.tensor([0] + [len(w) for w in merged_weight_distributions], dtype=torch.int64).cumsum(dim=0) + self.weight_distribution = torch.cat(merged_weight_distributions, dim=0).to(self.device) if self.bias: # Merge the bias distribution of all graphs (creating additionally slicing information) - self.bias_distribution_slices = torch.tensor([0] + [len(b) for b in self.bias_distribution], dtype=torch.int64).cumsum(dim=0) - self.bias_distribution = torch.cat([self.bias_distribution[i] for i in range(len(graph_data))], dim=0).to(self.device) + merged_bias_distributions = [ + torch.cat(chunks, dim=0) if chunks else torch.zeros((0, 4), dtype=torch.int64) + for chunks in bias_distribution_chunks + ] + self.bias_distribution_slices = torch.tensor([0] + [len(b) for b in merged_bias_distributions], dtype=torch.int64).cumsum(dim=0) + self.bias_distribution = torch.cat(merged_bias_distributions, dim=0).to(self.device) if self.bias: @@ -187,6 +188,10 @@ def __init__(self, parameters: Parameters, layer: Layer, graph_data: GraphDatase self.forward_step_time = 0 + # Cache config lookups used in forward pass + self.use_degree_matrix = self.para.run_config.config.get('degree_matrix', False) + self.use_in_degrees = self.para.run_config.config.get('use_in_degrees', False) + def init_weights(self, num_weights:np.float64, init_type:Optional[str]=None) -> nn.Parameter: """ Initializes the weights, i.e., learnable parameters of the module @@ -236,7 +241,7 @@ def set_weights(self, pos:int) -> None: :return: """ input_size = self.graph_data.num_nodes[pos].item() - self.current_W = torch.zeros((self.num_heads, input_size, input_size), dtype=self.precision).to(self.device) + self.current_W = torch.zeros((self.num_heads, input_size, input_size), dtype=self.precision, device=self.device) graph_weight_distribution = self.weight_distribution[self.weight_distribution_slices[pos]:self.weight_distribution_slices[pos+1]] if len(graph_weight_distribution) != 0: # get third column of the weight_distribution: the index of self.Param_W @@ -253,7 +258,7 @@ def set_bias(self, pos) -> None: :return: """ input_size = self.graph_data.num_nodes[pos].item() - self.current_B = torch.zeros((self.num_heads, input_size, self.in_features), dtype=self.precision).to(self.device) + self.current_B = torch.zeros((self.num_heads, input_size, self.in_features), dtype=self.precision, device=self.device) graph_bias_distribution = self.bias_distribution[self.bias_distribution_slices[pos]:self.bias_distribution_slices[pos+1]] param_indices = graph_bias_distribution[:, 3] matrix_indices = graph_bias_distribution[:, 0:3].T @@ -303,9 +308,9 @@ def forward(self, node_representation:torch.Tensor, batch_data: GraphDataset, *a begin = time.time() # set the weights, i.e., sets self.current_W to (C, N, N) where C is the number of channels and N is the number of nodes in graph at position pos of the dataset self.set_weights(pos) - if self.para.run_config.config.get('degree_matrix', False): + if self.use_degree_matrix: node_representation = self.in_edges[pos]*torch.einsum('cij,jk->cik', torch.diag(self.D[pos]) @ self.current_W @ torch.diag(self.D[pos]), node_representation) - elif self.para.run_config.config.get('use_in_degrees', False): + elif self.use_in_degrees: node_representation = self.in_edges[pos]*torch.einsum('cij,jk->cik', self.current_W, node_representation) else: node_representation = torch.einsum('hij,jf->hif', self.current_W, node_representation) diff --git a/src/models/ShareGNN/layers/inv_based_pooling.py b/src/models/ShareGNN/layers/inv_based_pooling.py index 710e738..5a33801 100644 --- a/src/models/ShareGNN/layers/inv_based_pooling.py +++ b/src/models/ShareGNN/layers/inv_based_pooling.py @@ -5,7 +5,6 @@ import networkx as nx import numpy as np import torch -from pandas.core.array_algos.masked_accumulations import cumsum from torch import nn from datasets.graph_dataset import GraphDataset @@ -36,36 +35,44 @@ def __init__(self, parameters:Parameters, layer: Layer, graph_data: GraphDataset self.n_node_labels = [] # number of node labels per head self.node_label_descriptions = [] # node label descriptions per head - self.n_heads_per_label = [] # number of heads per node label description # bias per head self.bias_list = [head.bias for head in layer.layer_heads] # is there any bias self.bias = any(self.bias_list) - for head_id, head in enumerate(layer.layer_heads): - self.node_label_descriptions.append(layer.get_source_string(head_id)) - self.n_node_labels.append(self.graph_data.node_labels[self.node_label_descriptions[head_id]].num_unique_node_labels) - self.n_heads_per_label.append(head.num) - - self.weight_num = np.sum([self.n_node_labels[i] * self.n_heads_per_label[i] for i in range(len(layer.layer_heads))]) - self.weight_distribution = torch.zeros((len(self.graph_data.x), self.num_heads), dtype=torch.int64) - # merge the bias distribution of all graphs (creating additionally slicing information) - self.weight_distribution_slices = self.graph_data.slices['x'] - - for head_id, head in enumerate(layer.layer_heads): - node_labels = self.graph_data.node_labels[self.node_label_descriptions[head_id]].node_labels + for i, head in enumerate(layer.layer_heads): + self.node_label_descriptions.append(layer.get_source_string(i)) + self.n_node_labels.append(self.graph_data.node_labels[self.node_label_descriptions[i]].num_unique_node_labels) + + self.weight_num = np.sum(self.n_node_labels) * self.output_dimension + weight_distribution_chunks = [[] for _ in range(len(self.graph_data))] + + for i, head in enumerate(layer.layer_heads): + node_labels = self.graph_data.node_labels[self.node_label_descriptions[i]].node_labels # Set the bias weights _, indices, counts = torch.unique(node_labels, dim=0, return_inverse=True, return_counts=True, sorted=False) for idx in range(len(self.graph_data)): - for h_num in range(head.num): - self.weight_distribution[self.graph_data.slices['x'][idx]:self.graph_data.slices['x'][idx+1], np.sum(self.n_heads_per_label[:head_id], dtype=int) + h_num] = indices[self.graph_data.slices['x'][idx]:self.graph_data.slices['x'][idx+1]] + h_num * self.n_node_labels[head_id] + np.sum([self.n_node_labels[i] * self.n_heads_per_label[i] for i in range(head_id)], dtype=int) - - + for out_dim_id in range(self.output_dimension): + new_weight_distribution = torch.zeros((self.graph_data.num_nodes[idx].item(), 4), dtype=torch.int64) + new_weight_distribution[:, 0] = i + new_weight_distribution[:, 1] = out_dim_id + new_weight_distribution[:, 2] = torch.arange(self.graph_data.num_nodes[idx].item()) # torch.arange(start=self.graph_data.slices['x'][idx], end=self.graph_data.slices['x'][idx+1], dtype=torch.int64) + new_weight_distribution[:, 3] = indices[self.graph_data.slices['x'][idx]:self.graph_data.slices['x'][idx+1]] + out_dim_id * self.n_node_labels[i] + weight_distribution_chunks[idx].append(new_weight_distribution) + + + # merge the weight distribution of all graphs (creating additionally slicing information) + merged_weight_distributions = [ + torch.cat(chunks, dim=0) if chunks else torch.zeros((0, 4), dtype=torch.int64) + for chunks in weight_distribution_chunks + ] + self.weight_distribution_slices = torch.tensor([0] + [len(w) for w in merged_weight_distributions], dtype=torch.int64).cumsum(dim=0) + self.weight_distribution = torch.cat(merged_weight_distributions, dim=0).to(self.device) self.Param_W = self.init_weights(self.weight_num, init_type='aggregation') if self.bias: - self.Param_b = self.init_weights(shape=(self.num_heads, self.in_features), init_type='aggregation_bias').to(self.device) + self.Param_b = self.init_weights(shape=(self.num_heads, self.output_dimension, self.in_features), init_type='aggregation_bias').to(self.device) self.forward_step_time = 0 @@ -116,9 +123,11 @@ def init_weights(self, shape, init_type=None): def set_weights(self, pos): input_size = self.graph_data.num_nodes[pos] - self.current_W = torch.zeros((input_size, self.num_heads), dtype=self.precision).to(self.device) + self.current_W = torch.zeros((self.num_heads, self.output_dimension, input_size), dtype=self.precision, device=self.device) weight_distr = self.weight_distribution[self.weight_distribution_slices[pos]:self.weight_distribution_slices[pos+1]] - self.current_W = torch.take(self.Param_W, weight_distr) + param_indices = weight_distr[:, 3] + matrix_indices = weight_distr[:, 0:3].T + self.current_W[matrix_indices[0], matrix_indices[1], matrix_indices[2]] = torch.take(self.Param_W, param_indices) # divide the weights by the number of nodes in the graph #self.current_W = self.current_W / input_size pass @@ -156,7 +165,7 @@ def forward(self, node_representation:torch.Tensor, batch_data: GraphDataset, *a pos = kwargs.get('pos', 0) begin = time.time() self.set_weights(pos) - node_representation = torch.einsum('no,nf->of', self.current_W, node_representation) + node_representation = torch.einsum('hoj,jf->hof', self.current_W, node_representation) if self.bias: node_representation = node_representation + self.Param_b node_representation = node_representation.flatten().unsqueeze(0) diff --git a/src/models/ShareGNN/layers/share_gnn_linear.py b/src/models/ShareGNN/layers/share_gnn_linear.py index 7f674a3..11500cd 100644 --- a/src/models/ShareGNN/layers/share_gnn_linear.py +++ b/src/models/ShareGNN/layers/share_gnn_linear.py @@ -30,9 +30,6 @@ def __init__(self, seed:int, layer_id:int, layer:Layer, parameters, graph_data:G self.in_features = in_features self.in_features = layer.layer_dict.get('in_features', self.in_features) - self.out_features = graph_data.num_node_features - if out_features is not None: - self.out_features = out_features self.out_features = layer.layer_dict.get('out_features', self.in_features) # determine the number of heads diff --git a/src/models/ShareGNN/preprocessing/preprocessing.py b/src/models/ShareGNN/preprocessing/preprocessing.py index 7f73ede..6b7f367 100644 --- a/src/models/ShareGNN/preprocessing/preprocessing.py +++ b/src/models/ShareGNN/preprocessing/preprocessing.py @@ -150,7 +150,8 @@ def layer_to_labels(experiment_configuration, layer_strings: json, graph_data: G raise ValueError( f'Please specify the subgraphs in the config files under the key "subgraphs" as folllows: subgraphs: - "[nx.complete_graph(4)]"') else: - subgraph_list = eval(experiment_configuration['subgraphs'][layer['id']]) + import ast + subgraph_list = ast.literal_eval(experiment_configuration['subgraphs'][layer['id']]) file_path = save_subgraph_labels(graph_data=graph_data, subgraphs=subgraph_list, subgraph_id=layer['id'], diff --git a/src/models/ShareGNN/utils.py b/src/models/ShareGNN/utils.py index 99ed1bc..00193dc 100644 --- a/src/models/ShareGNN/utils.py +++ b/src/models/ShareGNN/utils.py @@ -71,7 +71,6 @@ class LayerHead: """ def __init__(self, info_dict: dict, head_id): self.head_id = head_id - self.num = info_dict.get('num', 1) # How often the same head is applied self.label_dict = LabelDict(info_dict.get('labels', None)) self.property_dict = PropertyDict(info_dict.get('properties', None)) # if head, tail, bias is not specified, set head tail bias to the same value @@ -132,7 +131,7 @@ def get_unique_property_dicts(self) -> list[dict]: unique_dicts = [] for head in self.layer_heads: if head.property_dict.property_dict not in unique_dicts and head.property_dict.property_dict is not None: - unique_dicts.append(head.property_dict) + unique_dicts.append(head.property_dict.property_dict) return unique_dicts def get_source_string(self, head_id=0): diff --git a/src/models/layers/framework_layer.py b/src/models/layers/framework_layer.py index 5ccc9ec..5cdd93b 100644 --- a/src/models/layers/framework_layer.py +++ b/src/models/layers/framework_layer.py @@ -61,10 +61,6 @@ def __init__(self, layer_args=None): raise ValueError("out_channels must be provided") self.num_heads = layer_args.get('num_heads', 1) - if 'heads' in layer_args: - self.num_heads = 0 - for head in layer_args['heads']: - self.num_heads += head['num'] # Whether to use residual connections in this layer @@ -92,7 +88,28 @@ def __init__(self, layer_args=None): if 'activation' not in layer_args: self.activation = torch.nn.Identity() else: - self.activation = eval(layer_args['activation']) + activation_value = layer_args['activation'] + if isinstance(activation_value, torch.nn.Module): + self.activation = activation_value + elif isinstance(activation_value, str): + ACTIVATION_MAP = { + 'torch.nn.Identity()': torch.nn.Identity(), + 'torch.nn.ReLU()': torch.nn.ReLU(), + 'torch.nn.LeakyReLU()': torch.nn.LeakyReLU(), + 'torch.nn.ELU()': torch.nn.ELU(), + 'torch.nn.GELU()': torch.nn.GELU(), + 'torch.nn.Tanh()': torch.nn.Tanh(), + 'torch.nn.Sigmoid()': torch.nn.Sigmoid(), + 'torch.nn.Softmax(dim=1)': torch.nn.Softmax(dim=1), + 'torch.nn.SELU()': torch.nn.SELU(), + } + if activation_value in ACTIVATION_MAP: + self.activation = ACTIVATION_MAP[activation_value] + else: + raise ValueError(f"Unsupported activation function: '{activation_value}'. " + f"Supported values: {list(ACTIVATION_MAP.keys())}") + else: + raise ValueError(f"Activation must be a string or torch.nn.Module, got {type(activation_value)}") @abstractmethod def forward(self, node_representation:torch.Tensor, batch_data: GraphDataset, *args, **kwargs): diff --git a/src/models/layers/mpnn_classical/gat_conv.py b/src/models/layers/mpnn_classical/gat_conv.py index 87adb0c..4b90ef5 100644 --- a/src/models/layers/mpnn_classical/gat_conv.py +++ b/src/models/layers/mpnn_classical/gat_conv.py @@ -41,5 +41,5 @@ def forward(self, node_representation:torch.Tensor, batch_data: GraphDataset, *a else: # add residual to each head separately node_representation = node_representation + x.repeat(1, self.gat_args['heads']) if self.dropout > 0 and self.training: - node_representation = torch.nn.Dropout(self.dropout)(node_representation) + node_representation = torch.nn.functional.dropout(node_representation, p=self.dropout, training=self.training) return node_representation \ No newline at end of file diff --git a/src/models/layers/mpnn_classical/gatv2_conv.py b/src/models/layers/mpnn_classical/gatv2_conv.py index 89d66a5..786ae0d 100644 --- a/src/models/layers/mpnn_classical/gatv2_conv.py +++ b/src/models/layers/mpnn_classical/gatv2_conv.py @@ -46,5 +46,5 @@ def forward(self, node_representation:torch.Tensor, batch_data: GraphDataset, *a else: # add residual to each head separately node_representation = node_representation + x.repeat(1, self.gatv2_args['heads']) if self.dropout > 0 and self.training: - node_representation = torch.nn.Dropout(self.dropout)(node_representation) + node_representation = torch.nn.functional.dropout(node_representation, p=self.dropout, training=self.training) return node_representation \ No newline at end of file diff --git a/src/models/layers/mpnn_classical/gcn_conv.py b/src/models/layers/mpnn_classical/gcn_conv.py index f12e4b4..5b515b3 100644 --- a/src/models/layers/mpnn_classical/gcn_conv.py +++ b/src/models/layers/mpnn_classical/gcn_conv.py @@ -28,5 +28,5 @@ def forward(self,node_representation:torch.Tensor, batch_data: GraphDataset, *ar if self.residual: node_representation = node_representation + x if self.dropout > 0 and self.training: - node_representation = torch.nn.Dropout(self.dropout)(node_representation) + node_representation = torch.nn.functional.dropout(node_representation, p=self.dropout, training=self.training) return node_representation \ No newline at end of file diff --git a/src/models/layers/mpnn_classical/gin_conv.py b/src/models/layers/mpnn_classical/gin_conv.py index 2d65c7c..a096079 100644 --- a/src/models/layers/mpnn_classical/gin_conv.py +++ b/src/models/layers/mpnn_classical/gin_conv.py @@ -38,5 +38,5 @@ def forward(self, node_representation:torch.Tensor, batch_data: GraphDataset, *a if self.residual: node_representation = node_representation + x if self.dropout > 0 and self.training: - node_representation = torch.nn.Dropout(self.dropout)(node_representation) + node_representation = torch.nn.functional.dropout(node_representation, p=self.dropout, training=self.training) return node_representation \ No newline at end of file diff --git a/src/models/layers/mpnn_classical/sage_conv.py b/src/models/layers/mpnn_classical/sage_conv.py index f04bab1..8e84a94 100644 --- a/src/models/layers/mpnn_classical/sage_conv.py +++ b/src/models/layers/mpnn_classical/sage_conv.py @@ -29,5 +29,5 @@ def forward(self, node_representation:torch.Tensor, batch_data: GraphDataset, *a if self.residual: node_representation = node_representation + x if self.dropout > 0 and self.training: - node_representation = torch.nn.Dropout(self.dropout)(node_representation) + node_representation = torch.nn.functional.dropout(node_representation, p=self.dropout, training=self.training) return node_representation \ No newline at end of file diff --git a/src/models/layers/utils/layer_loader.py b/src/models/layers/utils/layer_loader.py index 1059f96..15de256 100644 --- a/src/models/layers/utils/layer_loader.py +++ b/src/models/layers/utils/layer_loader.py @@ -208,44 +208,44 @@ def check_layer(i:int, layer: dict)->(bool, str): if not isinstance(layer['heads'], list): return False, f'Heads must be a list in layer {i}' else: - for j, head in enumerate(layer['heads']): - if not isinstance(head, dict): + for j, heads in enumerate(layer['heads']): + if not isinstance(heads, dict): return False, f'Channel must be a dictionary in layer {i}' - if 'bias' not in head: + if 'bias' not in heads: return False, f'Bias not defined in layer {i}, it must be True or False' - if 'num' not in head: + if 'num' not in heads: return False, f'Number of heads must be defined in head {j}' - if 'labels' not in head: + if 'labels' not in heads: return False, f'Labels not defined in head {j}' else: if layer['layer_type'] == 'invariant_based_convolution': - if 'head' not in head['labels']: + if 'head' not in heads['labels']: return False, f'Head not defined in head {i}' else: - if 'label_type' not in head['labels']['head']: + if 'label_type' not in heads['labels']['head']: return False, f'Label type not defined in head {i} for head' - if 'tail' not in head['labels']: + if 'tail' not in heads['labels']: return False, f'Tail not defined in head {i}' else: - if 'label_type' not in head['labels']['tail']: + if 'label_type' not in heads['labels']['tail']: return False, f'Label type not defined in head {i} for tail' - if 'bias' not in head['labels']: + if 'bias' not in heads['labels']: return False, f'Bias not defined in head {i}' else: - if 'label_type' not in head['labels']['bias']: + if 'label_type' not in heads['labels']['bias']: return False, f'Label type not defined in head {i} for bias' - if 'properties' not in head: + if 'properties' not in heads: return False, f'Properties not defined in head {i}' else: - if 'name' not in head['properties']: + if 'name' not in heads['properties']: return False, f'Property name not defined in head {i}' - if 'values' not in head['properties']: + if 'values' not in heads['properties']: return False, f'Property values not defined in head {i}' else: - if not isinstance(head['properties']['values'], list): + if not isinstance(heads['properties']['values'], list): return False, f'Property values must be a list in head {i}' elif layer['layer_type'] == 'invariant_based_aggregation': - if 'label_type' not in head['labels']: + if 'label_type' not in heads['labels']: return False, f'Label type not defined in head {i}' else: return False, f'Layer type {layer["layer_type"]} not supported in layer {i}' diff --git a/src/models/model.py b/src/models/model.py index 600f28a..bcb7082 100644 --- a/src/models/model.py +++ b/src/models/model.py @@ -69,14 +69,8 @@ def forward(self, batch_data, *args, **kwargs): dtype=torch.double) x = x + random_variation - representation_list = [] - for i, layer in enumerate(self.net_layers): + for layer in self.net_layers: x = layer(x, batch_data, *args, **kwargs) - continue - #if isinstance(layer, GNNConvLayer): - # representation_list.append(x) - # if i == len(self.net_layers) - 1 or not isinstance(self.net_layers[i + 1], GNNConvLayer): - # x = torch.squeeze(torch.mean(torch.stack(representation_list), 0, True), 0) return x @@ -130,13 +124,13 @@ def get_model_layer(self, layer): layer_args = layer.layer_dict # GNN specific layers if layer.layer_type == LayerTypes.GCN_CONVOLUTION.value: - self.net_layers.append(GCNConv(layer_args).type(self.precision).requires_grad_(self.convolution_grad)) + return GCNConv(layer_args).type(self.precision).requires_grad_(self.convolution_grad) elif layer.layer_type == LayerTypes.GAT_CONVOLUTION.value: - self.net_layers.append(GATConv(layer_args).type(self.precision).requires_grad_(self.convolution_grad)) + return GATConv(layer_args).type(self.precision).requires_grad_(self.convolution_grad) elif layer.layer_type == LayerTypes.GATv2_CONVOLUTION.value: - self.net_layers.append(GATv2Conv(layer_args).type(self.precision).requires_grad_(self.convolution_grad)) + return GATv2Conv(layer_args).type(self.precision).requires_grad_(self.convolution_grad) elif layer.layer_type == LayerTypes.GIN_CONVOLUTION.value: - self.net_layers.append(GINConv(layer_args).type(self.precision).requires_grad_(self.convolution_grad)) + return GINConv(layer_args).type(self.precision).requires_grad_(self.convolution_grad) elif layer.layer_type == LayerTypes.SAGE_CONVOLUTION.value: return SAGEConv(layer_args).type(self.precision).requires_grad_(self.convolution_grad) From 14e242fc19fd94804423389560ecd13828a27f6e Mon Sep 17 00:00:00 2001 From: florian Date: Mon, 9 Feb 2026 12:32:28 +0100 Subject: [PATCH 3/4] Refactor head handling in layers to support multiple heads per label and improve weight distribution logic TODO fix inv_based_message_passing.py --- .claude/CLAUDE.md | 131 ----------------------------------- src/models/ShareGNN/utils.py | 2 +- 2 files changed, 1 insertion(+), 132 deletions(-) delete mode 100644 .claude/CLAUDE.md diff --git a/.claude/CLAUDE.md b/.claude/CLAUDE.md deleted file mode 100644 index c9b931a..0000000 --- a/.claude/CLAUDE.md +++ /dev/null @@ -1,131 +0,0 @@ -# CLAUDE.md - SimpleGNN Project Guide - -## Project Overview - -SimpleGNN is a PyTorch-based Graph Neural Network experimentation framework for benchmarking and developing GNN architectures. It supports classical message-passing GNNs (GCN, GIN, GAT, GATv2, GraphSAGE) and a proprietary ShareGNN variant with invariant-based layers. The framework handles graph classification, graph regression, and node classification tasks. - -## Repository Structure - -``` -SimpleGNN/ -├── src/ # All source code -│ ├── framework/ # Training/evaluation orchestration -│ │ ├── core.py # FrameworkMain - main entry point class -│ │ ├── model_configuration.py # Single config training: model init, train/eval loops -│ │ ├── run_configuration.py # Hyperparameter grid search combinations -│ │ └── utils/ -│ │ ├── parameters.py # Parameters class (all experiment hyperparams) -│ │ ├── preprocessing.py # Dataset loading, splits, label computation -│ │ ├── evaluation.py # Result analysis and visualization -│ │ ├── configuration_checks.py # YAML config validation -│ │ └── data_sampling.py # Batch sampling strategies -│ ├── models/ # GNN models and layers -│ │ ├── model.py # GraphModel - main PyTorch model class -│ │ ├── layers/ -│ │ │ ├── framework_layer.py # Abstract base class for all custom layers -│ │ │ ├── mpnn_classical/ # Classical GNN wrappers (gcn, gin, gat, gatv2, sage, pooling) -│ │ │ ├── nn_standard/ # Standard layers (linear, activation, batchnorm, dropout, reshape) -│ │ │ └── utils/ # LayerTypes enum, layer_loader -│ │ └── ShareGNN/ # Proprietary ShareGNN implementation -│ │ ├── layers/ # inv_based_message_passing, inv_based_pooling, positional_encoding -│ │ ├── preprocessing/ # ShareGNN-specific label/property preprocessing -│ │ └── utils.py -│ ├── datasets/ # Dataset handling -│ │ ├── graph_dataset.py # GraphDataset (PyG InMemoryDataset) -│ │ ├── graph_dataset_preprocessing.py # Dataset-specific preprocessing -│ │ ├── custom_datasets.py # Dataset factory and registration -│ │ ├── custom_benchmarks/ # Synthetic benchmarks (rings, snowflakes, strings) -│ │ ├── splits/ # Pre-computed train/val/test splits (JSON) -│ │ └── utils/ # Node/edge label utilities, graph functions -│ └── utils/ # General utilities (timer, path conversions) -├── examples/ # Example experiments with YAML configs -│ ├── basic_example/ # GCN/GIN on MUTAG -│ ├── basic_example_share_gnn/ # ShareGNN on MUTAG -│ └── zinc/ # ShareGNN on ZINC -├── experiments/ # Shell scripts for reproducing paper results -│ └── base_paper/ # TUDatasets, ZINC, QM9, ablation, synthetic benchmarks -├── data/ # Dataset storage (gitignored, auto-downloaded) -├── results/ # Experiment results (gitignored) -└── docs/ # Sphinx documentation -``` - -## Configuration System - -The framework uses a **three-tier YAML configuration**: - -1. **Main config** (`main.yml`): Datasets, task type, paths to model/hyperparameter configs, splits -2. **Model config** (`models_*.yml`): Layer architecture as list of layer definitions (supports grid search via list of lists) -3. **Hyperparameter config** (`parameters.yml`): Training params (optimizer, loss, lr, epochs, batch size, input features). Lists = grid search. - -Configs are validated by `src/framework/utils/configuration_checks.py` against mandatory parameter sets. - -## Running Experiments - -Entry point pattern (see `examples/*/main.py`): - -```python -from pathlib import Path -from framework.core import FrameworkMain - -experiment = FrameworkMain(Path('examples/basic_example/main.yml')) -experiment.preprocessing(num_threads=1) # Load data, generate labels/splits -experiment.run_configurations(num_threads=-1) # Grid search (-1 = all CPUs) -experiment.evaluate_results() # Find best config on validation set -experiment.run_best_configuration(num_threads=-1) # Re-run best config -experiment.evaluate_results(evaluate_best_model=True) # Final test-set evaluation -``` - -Run from the `src/` directory (imports are relative to `src/`): -```bash -cd src && python -m examples.basic_example.main -``` - -Experiment shell scripts are in `experiments/base_paper/`. - -## Key Architecture Decisions - -- **GraphModel** (`src/models/model.py`): Sequential `nn.ModuleList` of layers built from YAML config. All layers extend `FrameworkLayer` base class. -- **Layer types**: Defined in `LayerTypes` enum (`src/models/layers/utils/layer_types.py`). New layers must be registered there and in `layer_loader.py`. -- **Tensor shapes**: Layers handle `(C, N, F)` for multi-channel/multi-head data or `(N, F)` for standard. C = channels/heads, N = nodes, F = features. -- **LinearLayer modes**: `aggr_features` (standard), `aggr_channels` (aggregate across channels), `channel_wise` (independent per channel). -- **ShareGNN layers** use node/edge labels and pairwise properties for invariant-based message passing (multi-head output). - -## Dependencies - -Core: PyTorch, PyTorch Geometric (torch_geometric), numpy, pandas, scikit-learn, networkx, joblib, pyyaml, matplotlib. -Optional: ogb, rdkit (molecular datasets). - -Virtual environment is in `venv/`. Note: `requirements.txt` and `setup.py` have been removed from tracking. - -## Coding Conventions - -- **Classes**: PascalCase (`FrameworkMain`, `GraphDataset`, `GCNConv`) -- **Functions/methods**: snake_case (`run_configurations`, `forward`, `load_preprocessed_data_and_parameters`) -- **Constants**: UPPER_SNAKE_CASE (`MANDATORY_MAIN_CONFIG_PARAMS`) -- **Private members**: underscore prefix (`_num_graph_nodes`) -- **Imports**: Standard library first, then third-party (torch, numpy), then local modules. Local imports use dot-notation relative to `src/` (e.g., `from models.model import GraphModel`). -- **Type hints**: Used in method signatures where present, not comprehensive across the codebase. -- **Indentation**: 4 spaces. -- **Docstrings**: Present on classes and key methods; not exhaustive. Don't add docstrings to code you didn't write. - -## Git Conventions - -- **Main branch**: `master` -- **Commit messages**: Imperative mood, descriptive ("Refactor linear layer implementation to support 'aggr_channels' mode") -- **No unit test suite**: Testing is done via integration experiments (example scripts and experiment shell scripts) -- **Gitignored**: `data/`, `results/`, `*.csv`, `venv/`, `__pycache__/`, `.idea/`, `.auto-claude/`, `/.claude/` - -## Important Patterns - -- **Adding a new GNN layer**: Create wrapper class extending `FrameworkLayer` in `src/models/layers/mpnn_classical/`, add to `LayerTypes` enum, register in `layer_loader.py`, then reference in YAML model config. -- **Adding a new dataset**: Implement preprocessing class in `graph_dataset_preprocessing.py`, register in `custom_datasets.py`, create split files in `datasets/splits/`. -- **Grid search**: Use lists in YAML parameter files. The framework computes the cartesian product of all list-valued parameters. -- **Parallel execution**: `joblib` handles parallel runs across splits/configs. `num_threads=-1` uses all CPUs. -- **Results**: Saved as CSV per epoch per configuration in the results directory. Evaluation finds best config by validation metric. - -## Common Pitfalls - -- Always run from `src/` directory or ensure `src/` is on the Python path; imports are relative to it. -- YAML model configs use `layer_type` string keys that must match `LayerTypes` enum values exactly. -- ShareGNN preprocessing must run before training ShareGNN models (generates required node/edge properties). -- The `precision` parameter (`float`/`double`) must be consistent between data and model; mismatches cause runtime errors. diff --git a/src/models/ShareGNN/utils.py b/src/models/ShareGNN/utils.py index 00193dc..21aaa84 100644 --- a/src/models/ShareGNN/utils.py +++ b/src/models/ShareGNN/utils.py @@ -131,7 +131,7 @@ def get_unique_property_dicts(self) -> list[dict]: unique_dicts = [] for head in self.layer_heads: if head.property_dict.property_dict not in unique_dicts and head.property_dict.property_dict is not None: - unique_dicts.append(head.property_dict.property_dict) + unique_dicts.append(head.property_dict) return unique_dicts def get_source_string(self, head_id=0): From 397906b14e0bd7a9ef8618f42609dd95f39ba558 Mon Sep 17 00:00:00 2001 From: florian Date: Mon, 9 Feb 2026 14:09:24 +0100 Subject: [PATCH 4/4] Refactor layer handling to support multiple heads per label and improve weight distribution logic --- src/datasets/graph_dataset.py | 4 +- src/framework/model_configuration.py | 13 ++- src/framework/utils/evaluation.py | 2 +- .../layers/inv_based_message_passing.py | 88 ++++++++++++------- .../ShareGNN/layers/inv_based_pooling.py | 53 +++++------ .../ShareGNN/preprocessing/properties.py | 2 +- src/models/ShareGNN/utils.py | 1 + src/models/layers/framework_layer.py | 4 + src/models/layers/utils/layer_loader.py | 32 +++---- src/models/model.py | 3 +- 10 files changed, 109 insertions(+), 93 deletions(-) diff --git a/src/datasets/graph_dataset.py b/src/datasets/graph_dataset.py index 3f0f3b9..81cc724 100644 --- a/src/datasets/graph_dataset.py +++ b/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 diff --git a/src/framework/model_configuration.py b/src/framework/model_configuration.py index e3db69e..3ba44a9 100644 --- a/src/framework/model_configuration.py +++ b/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 @@ -60,7 +60,7 @@ def __init__(self, run_id: int, k_val: int, graph_data: GraphDataset, model_data self.net = None self.class_weights = None self._csv_buffer = [] - self._csv_flush_interval = 50 + 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") @@ -353,8 +353,7 @@ def evaluate_results(self, epoch: int, batch_length=0, num_batches=0, batches=None): - if evaluation_type != 'training': - self.net.eval() + self.net.eval() if evaluation_type == 'training': batch_acc = 0 # if num classes is one calculate the mae and mae_std or if the task is regression @@ -639,8 +638,7 @@ def evaluate_results(self, epoch: int, labels_np = labels_np.T df = pd.DataFrame(labels_np) df.to_csv("Results/Parameter/test_predictions.csv", header=False, index=False, mode='a') - if evaluation_type != 'training': - self.net.train() + self.net.train() return train_values, validation_values, test_values def preprocess_writer(self)-> bool: @@ -787,7 +785,6 @@ def postprocess_writer(self, epoch, epoch_time, train_values: EvaluationValues, 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 @@ -872,7 +869,6 @@ def evaluate_graph_task(self, graph_ids): with torch.no_grad(): self.net.train(False) - # Run ordinary GNN # 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)] @@ -887,6 +883,7 @@ 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) diff --git a/src/framework/utils/evaluation.py b/src/framework/utils/evaluation.py index 32dfe58..b66a95d 100644 --- a/src/framework/utils/evaluation.py +++ b/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): diff --git a/src/models/ShareGNN/layers/inv_based_message_passing.py b/src/models/ShareGNN/layers/inv_based_message_passing.py index a836f75..1acbfa9 100644 --- a/src/models/ShareGNN/layers/inv_based_message_passing.py +++ b/src/models/ShareGNN/layers/inv_based_message_passing.py @@ -47,6 +47,7 @@ def __init__(self, parameters: Parameters, layer: Layer, graph_data: GraphDatase super(InvariantBasedMessagePassingLayer, self).__init__(parameters, layer, graph_data) self.out_features = self.in_features * self.num_heads + self.n_heads_per_label = [] # number of heads per node label description for h_id, head in enumerate(layer.layer_heads): self.source_label_descriptions.append(layer.get_source_string(h_id)) @@ -57,30 +58,38 @@ def __init__(self, parameters: Parameters, layer: Layer, graph_data: GraphDatase self.n_bias_labels.append(graph_data.node_labels[self.bias_label_descriptions[h_id]].num_unique_node_labels) self.property_descriptions.append(head.property_dict.get_property_string()) self.n_properties.append(graph_data.properties[self.property_descriptions[h_id]].num_properties[(self.layer_id, h_id)]) + self.n_heads_per_label.append(head.num) self.bias_list = [head.bias for head in layer.layer_heads] self.bias = any(self.bias_list) # check if bias is used + + # Determine the number of weights and biases # There are two cases asymetric and symmetric, asymetric is the default, TODO add symmetric case self.skips = [0] + self.b_head_offset = 0 self.skips_description = [None] self.skips_description_text = [None] weight_distribution_chunks = [[] for _ in range(len(graph_data))] bias_distribution_chunks = [[] for _ in range(len(graph_data))] # Iterate over all heads in the layer - for i, head in enumerate(self.layer.layer_heads): + for head_id, head in enumerate(self.layer.layer_heads): # get all the valid property values for the head (e.g., the distances 0, 3, 6) - valid_property_values = self.graph_data.properties[self.property_descriptions[i]].valid_values[(self.layer_id, i)] + valid_property_values = self.graph_data.properties[self.property_descriptions[head_id]].valid_values[(self.layer_id, head_id)] # apply the head and tail labels to the subdict - source_labels = self.graph_data.node_labels[self.source_label_descriptions[i]].node_labels - target_labels = self.graph_data.node_labels[self.target_label_descriptions[i]].node_labels - bias_labels = self.graph_data.node_labels[self.bias_label_descriptions[i]].node_labels + source_labels = self.graph_data.node_labels[self.source_label_descriptions[head_id]].node_labels + target_labels = self.graph_data.node_labels[self.target_label_descriptions[head_id]].node_labels + bias_labels = self.graph_data.node_labels[self.bias_label_descriptions[head_id]].node_labels + current_head_id = 0 + for h_i in range(head_id): + current_head_id += self.n_heads_per_label[h_i] + for key in valid_property_values: #print(f'Initialize head {i+1}/{len(self.layer.layer_heads)} with property {key}') - property_subdict = self.graph_data.properties[self.property_descriptions[i]].properties[key] - property_subdict_slices = self.graph_data.properties[self.property_descriptions[i]].properties_slices[key] + property_subdict = self.graph_data.properties[self.property_descriptions[head_id]].properties[key] + property_subdict_slices = self.graph_data.properties[self.property_descriptions[head_id]].properties_slices[key] labeled_subdict = property_subdict.detach().clone() labeled_subdict[:, 0] = source_labels[property_subdict[:, 0]] labeled_subdict[:, 1] = target_labels[property_subdict[:, 1]] @@ -113,46 +122,59 @@ def __init__(self, parameters: Parameters, layer: Layer, graph_data: GraphDatase # relabel indices indices[valid_indices] = torch.tensor([valid_value_dict[idx.item()] for idx in indices[valid_indices]], dtype=torch.int64) num_weights = len(valid_values) + for n in range(self.n_heads_per_label[head_id]): + self.weight_num.append(num_weights) start_time = time.time() for idx in range(len(graph_data)): # if number of graphs is larger than 10000 print progress - if len(graph_data) > 10000 and idx % 1000 == 0: - print(f'Head {i+1}/{len(self.layer.layer_heads)} with property {key}: {idx}/{len(graph_data)} graphs processed ({(idx/len(graph_data))*100:.2f}%) time so far (in s): {time.time()-start_time:.2f}', flush=True) + if len(graph_data) > 5000 and idx % 1000 == 0: + print( + f'Heads {self.n_heads_per_label[head_id] + current_head_id}/{self.num_heads} with property {key}: {idx}/{len(graph_data)} graphs processed ({(idx / len(graph_data)) * 100:.2f}%) time so far (in s): {time.time() - start_time:.2f}', + flush=True) # get the valid indices for the current graph if threshold > 1 or do_invalid_indices_exist or upper_threshold is not None: - valid_indices_graph = torch.where(valid_indices_bool[property_subdict_slices[idx]:property_subdict_slices[idx+1]])[0] + property_subdict_slices[idx] + valid_indices_graph = torch.where( + valid_indices_bool[property_subdict_slices[idx]:property_subdict_slices[idx + 1]])[0] + \ + property_subdict_slices[idx] else: - valid_indices_graph = torch.arange(property_subdict_slices[idx], property_subdict_slices[idx+1], dtype=torch.int64) + valid_indices_graph = torch.arange(property_subdict_slices[idx], + property_subdict_slices[idx + 1], dtype=torch.int64) + w_indices = indices[valid_indices_graph] + p_indices = property_subdict[valid_indices_graph] - self.graph_data.slices['x'][idx] # check if subtracting is necessary # create new tensor where each row is the concatenation of head_id, property_subdict_row, and indices - new_weight_distribution = torch.zeros((len(valid_indices_graph), 4), dtype=torch.int64) - new_weight_distribution[:, 0] = i - new_weight_distribution[:, 1:3] = property_subdict[valid_indices_graph] - self.graph_data.slices['x'][idx] # check if subtracting is necessary - new_weight_distribution[:, 3] = indices[valid_indices_graph] + self.skips[-1] - weight_distribution_chunks[idx].append(new_weight_distribution) - - + for n in range(self.n_heads_per_label[head_id]): + new_weight_distribution = torch.zeros((len(valid_indices_graph), 4), dtype=torch.int64) + new_weight_distribution[:, 0] = current_head_id + n + new_weight_distribution[:, 1:3] = p_indices + new_weight_distribution[:, 3] = w_indices + self.skips[-1] + n*num_weights + weight_distribution_chunks[idx].append(new_weight_distribution) + + for n in range(self.n_heads_per_label[head_id]): + self.skips.append(self.skips[-1] + num_weights) + self.skips_description.append({'head:': head_id, 'property': key, 'weights': num_weights}) + self.skips_description_text.append(f"Head {head_id} Property {key} has {num_weights} different weights") - self.skips.append(self.skips[-1] + num_weights) - self.skips_description.append({'head:': i, 'property': key, 'weights': num_weights}) - self.skips_description_text.append(f"Head {i} Property {key} has {num_weights} different weights") - - - self.weight_num.append(self.skips[-1]) # TODO symmetric case if self.bias: - # Determine the number of different learnable parameters in the bias vector - self.bias_num.append(self.in_features * self.n_bias_labels[i]) # Set the bias weights _, indices, counts = torch.unique(bias_labels, dim=0, return_inverse=True, return_counts=True, sorted=False) for idx in range(len(graph_data)): - for feature_id in range(self.in_features): - new_bias_distribution = torch.zeros((graph_data.num_nodes[idx].item(), 4), dtype=torch.int64) - new_bias_distribution[:, 0] = i - new_bias_distribution[:, 1] = torch.arange(graph_data.num_nodes[idx].item(), dtype=torch.int64) # alternative torch.arange(start=graph_data.slices['x'][idx], end=graph_data.slices['x'][idx+1], dtype=torch.int64) - new_bias_distribution[:, 2] = feature_id - new_bias_distribution[:, 3] = indices[graph_data.slices['x'][idx]:graph_data.slices['x'][idx+1]] + feature_id * self.n_bias_labels[i] - bias_distribution_chunks[idx].append(new_bias_distribution) + arranged_tensor = torch.arange(graph_data.num_nodes[idx].item(), dtype=torch.int64) # alternative torch.arange(start=graph_data.slices['x'][idx], end=graph_data.slices['x'][idx+1], dtype=torch.int64) + w_indices = indices[graph_data.slices['x'][idx]:graph_data.slices['x'][idx+1]] + for n in range(self.n_heads_per_label[head_id]): + for feature_id in range(self.in_features): + w_index_offset = n*self.in_features*self.n_bias_labels[head_id] + feature_id*self.n_bias_labels[head_id] + new_bias_distribution = torch.zeros((graph_data.num_nodes[idx].item(), 4), dtype=torch.int64) + new_bias_distribution[:, 0] = current_head_id + n + new_bias_distribution[:, 1] = arranged_tensor + new_bias_distribution[:, 2] = feature_id + new_bias_distribution[:, 3] = w_indices + self.b_head_offset + w_index_offset + bias_distribution_chunks[idx].append(new_bias_distribution) + # Determine the number of different learnable parameters in the bias vector + for n in range(self.n_heads_per_label[head_id]): + self.bias_num.append(self.in_features * self.n_bias_labels[head_id]) + self.b_head_offset += self.bias_num[-1] # Merge the weight distribution of all graphs (creating additionally slicing information) # Single torch.cat per graph instead of repeated incremental concatenation diff --git a/src/models/ShareGNN/layers/inv_based_pooling.py b/src/models/ShareGNN/layers/inv_based_pooling.py index 5a33801..710e738 100644 --- a/src/models/ShareGNN/layers/inv_based_pooling.py +++ b/src/models/ShareGNN/layers/inv_based_pooling.py @@ -5,6 +5,7 @@ import networkx as nx import numpy as np import torch +from pandas.core.array_algos.masked_accumulations import cumsum from torch import nn from datasets.graph_dataset import GraphDataset @@ -35,44 +36,36 @@ def __init__(self, parameters:Parameters, layer: Layer, graph_data: GraphDataset self.n_node_labels = [] # number of node labels per head self.node_label_descriptions = [] # node label descriptions per head + self.n_heads_per_label = [] # number of heads per node label description # bias per head self.bias_list = [head.bias for head in layer.layer_heads] # is there any bias self.bias = any(self.bias_list) - for i, head in enumerate(layer.layer_heads): - self.node_label_descriptions.append(layer.get_source_string(i)) - self.n_node_labels.append(self.graph_data.node_labels[self.node_label_descriptions[i]].num_unique_node_labels) - - self.weight_num = np.sum(self.n_node_labels) * self.output_dimension - weight_distribution_chunks = [[] for _ in range(len(self.graph_data))] - - for i, head in enumerate(layer.layer_heads): - node_labels = self.graph_data.node_labels[self.node_label_descriptions[i]].node_labels + for head_id, head in enumerate(layer.layer_heads): + self.node_label_descriptions.append(layer.get_source_string(head_id)) + self.n_node_labels.append(self.graph_data.node_labels[self.node_label_descriptions[head_id]].num_unique_node_labels) + self.n_heads_per_label.append(head.num) + + self.weight_num = np.sum([self.n_node_labels[i] * self.n_heads_per_label[i] for i in range(len(layer.layer_heads))]) + self.weight_distribution = torch.zeros((len(self.graph_data.x), self.num_heads), dtype=torch.int64) + # merge the bias distribution of all graphs (creating additionally slicing information) + self.weight_distribution_slices = self.graph_data.slices['x'] + + for head_id, head in enumerate(layer.layer_heads): + node_labels = self.graph_data.node_labels[self.node_label_descriptions[head_id]].node_labels # Set the bias weights _, indices, counts = torch.unique(node_labels, dim=0, return_inverse=True, return_counts=True, sorted=False) for idx in range(len(self.graph_data)): - for out_dim_id in range(self.output_dimension): - new_weight_distribution = torch.zeros((self.graph_data.num_nodes[idx].item(), 4), dtype=torch.int64) - new_weight_distribution[:, 0] = i - new_weight_distribution[:, 1] = out_dim_id - new_weight_distribution[:, 2] = torch.arange(self.graph_data.num_nodes[idx].item()) # torch.arange(start=self.graph_data.slices['x'][idx], end=self.graph_data.slices['x'][idx+1], dtype=torch.int64) - new_weight_distribution[:, 3] = indices[self.graph_data.slices['x'][idx]:self.graph_data.slices['x'][idx+1]] + out_dim_id * self.n_node_labels[i] - weight_distribution_chunks[idx].append(new_weight_distribution) - - - # merge the weight distribution of all graphs (creating additionally slicing information) - merged_weight_distributions = [ - torch.cat(chunks, dim=0) if chunks else torch.zeros((0, 4), dtype=torch.int64) - for chunks in weight_distribution_chunks - ] - self.weight_distribution_slices = torch.tensor([0] + [len(w) for w in merged_weight_distributions], dtype=torch.int64).cumsum(dim=0) - self.weight_distribution = torch.cat(merged_weight_distributions, dim=0).to(self.device) + for h_num in range(head.num): + self.weight_distribution[self.graph_data.slices['x'][idx]:self.graph_data.slices['x'][idx+1], np.sum(self.n_heads_per_label[:head_id], dtype=int) + h_num] = indices[self.graph_data.slices['x'][idx]:self.graph_data.slices['x'][idx+1]] + h_num * self.n_node_labels[head_id] + np.sum([self.n_node_labels[i] * self.n_heads_per_label[i] for i in range(head_id)], dtype=int) + + self.Param_W = self.init_weights(self.weight_num, init_type='aggregation') if self.bias: - self.Param_b = self.init_weights(shape=(self.num_heads, self.output_dimension, self.in_features), init_type='aggregation_bias').to(self.device) + self.Param_b = self.init_weights(shape=(self.num_heads, self.in_features), init_type='aggregation_bias').to(self.device) self.forward_step_time = 0 @@ -123,11 +116,9 @@ def init_weights(self, shape, init_type=None): def set_weights(self, pos): input_size = self.graph_data.num_nodes[pos] - self.current_W = torch.zeros((self.num_heads, self.output_dimension, input_size), dtype=self.precision, device=self.device) + self.current_W = torch.zeros((input_size, self.num_heads), dtype=self.precision).to(self.device) weight_distr = self.weight_distribution[self.weight_distribution_slices[pos]:self.weight_distribution_slices[pos+1]] - param_indices = weight_distr[:, 3] - matrix_indices = weight_distr[:, 0:3].T - self.current_W[matrix_indices[0], matrix_indices[1], matrix_indices[2]] = torch.take(self.Param_W, param_indices) + self.current_W = torch.take(self.Param_W, weight_distr) # divide the weights by the number of nodes in the graph #self.current_W = self.current_W / input_size pass @@ -165,7 +156,7 @@ def forward(self, node_representation:torch.Tensor, batch_data: GraphDataset, *a pos = kwargs.get('pos', 0) begin = time.time() self.set_weights(pos) - node_representation = torch.einsum('hoj,jf->hof', self.current_W, node_representation) + node_representation = torch.einsum('no,nf->of', self.current_W, node_representation) if self.bias: node_representation = node_representation + self.Param_b node_representation = node_representation.flatten().unsqueeze(0) diff --git a/src/models/ShareGNN/preprocessing/properties.py b/src/models/ShareGNN/preprocessing/properties.py index 8973101..c35e7d8 100644 --- a/src/models/ShareGNN/preprocessing/properties.py +++ b/src/models/ShareGNN/preprocessing/properties.py @@ -11,7 +11,7 @@ from datasets.graph_dataset import GraphDataset from datasets.utils.node_labeling import load_labels -from src.utils.utils import convert_to_list +from utils.utils import convert_to_list def write_distance_properties(graph_data:GraphDataset, cutoff=None, out_path: Path = Path(), save_times=None) -> None: diff --git a/src/models/ShareGNN/utils.py b/src/models/ShareGNN/utils.py index 21aaa84..99ed1bc 100644 --- a/src/models/ShareGNN/utils.py +++ b/src/models/ShareGNN/utils.py @@ -71,6 +71,7 @@ class LayerHead: """ def __init__(self, info_dict: dict, head_id): self.head_id = head_id + self.num = info_dict.get('num', 1) # How often the same head is applied self.label_dict = LabelDict(info_dict.get('labels', None)) self.property_dict = PropertyDict(info_dict.get('properties', None)) # if head, tail, bias is not specified, set head tail bias to the same value diff --git a/src/models/layers/framework_layer.py b/src/models/layers/framework_layer.py index 5cdd93b..02fba87 100644 --- a/src/models/layers/framework_layer.py +++ b/src/models/layers/framework_layer.py @@ -61,6 +61,10 @@ def __init__(self, layer_args=None): raise ValueError("out_channels must be provided") self.num_heads = layer_args.get('num_heads', 1) + if 'heads' in layer_args: + self.num_heads = 0 + for head in layer_args['heads']: + self.num_heads += head['num'] # Whether to use residual connections in this layer diff --git a/src/models/layers/utils/layer_loader.py b/src/models/layers/utils/layer_loader.py index 15de256..1059f96 100644 --- a/src/models/layers/utils/layer_loader.py +++ b/src/models/layers/utils/layer_loader.py @@ -208,44 +208,44 @@ def check_layer(i:int, layer: dict)->(bool, str): if not isinstance(layer['heads'], list): return False, f'Heads must be a list in layer {i}' else: - for j, heads in enumerate(layer['heads']): - if not isinstance(heads, dict): + for j, head in enumerate(layer['heads']): + if not isinstance(head, dict): return False, f'Channel must be a dictionary in layer {i}' - if 'bias' not in heads: + if 'bias' not in head: return False, f'Bias not defined in layer {i}, it must be True or False' - if 'num' not in heads: + if 'num' not in head: return False, f'Number of heads must be defined in head {j}' - if 'labels' not in heads: + if 'labels' not in head: return False, f'Labels not defined in head {j}' else: if layer['layer_type'] == 'invariant_based_convolution': - if 'head' not in heads['labels']: + if 'head' not in head['labels']: return False, f'Head not defined in head {i}' else: - if 'label_type' not in heads['labels']['head']: + if 'label_type' not in head['labels']['head']: return False, f'Label type not defined in head {i} for head' - if 'tail' not in heads['labels']: + if 'tail' not in head['labels']: return False, f'Tail not defined in head {i}' else: - if 'label_type' not in heads['labels']['tail']: + if 'label_type' not in head['labels']['tail']: return False, f'Label type not defined in head {i} for tail' - if 'bias' not in heads['labels']: + if 'bias' not in head['labels']: return False, f'Bias not defined in head {i}' else: - if 'label_type' not in heads['labels']['bias']: + if 'label_type' not in head['labels']['bias']: return False, f'Label type not defined in head {i} for bias' - if 'properties' not in heads: + if 'properties' not in head: return False, f'Properties not defined in head {i}' else: - if 'name' not in heads['properties']: + if 'name' not in head['properties']: return False, f'Property name not defined in head {i}' - if 'values' not in heads['properties']: + if 'values' not in head['properties']: return False, f'Property values not defined in head {i}' else: - if not isinstance(heads['properties']['values'], list): + if not isinstance(head['properties']['values'], list): return False, f'Property values must be a list in head {i}' elif layer['layer_type'] == 'invariant_based_aggregation': - if 'label_type' not in heads['labels']: + if 'label_type' not in head['labels']: return False, f'Label type not defined in head {i}' else: return False, f'Layer type {layer["layer_type"]} not supported in layer {i}' diff --git a/src/models/model.py b/src/models/model.py index bcb7082..c095d9c 100644 --- a/src/models/model.py +++ b/src/models/model.py @@ -69,7 +69,7 @@ def forward(self, batch_data, *args, **kwargs): dtype=torch.double) x = x + random_variation - for layer in self.net_layers: + for i, layer in enumerate(self.net_layers): x = layer(x, batch_data, *args, **kwargs) return x @@ -134,6 +134,7 @@ def get_model_layer(self, layer): elif layer.layer_type == LayerTypes.SAGE_CONVOLUTION.value: return SAGEConv(layer_args).type(self.precision).requires_grad_(self.convolution_grad) + elif layer.layer_type == LayerTypes.GLOBAL_POOLING.value: layer_args = {'mode': layer.layer_dict.get('mode', 'mean')} return GlobalPooling(layer_args).type(self.precision).requires_grad_(self.aggregation_grad)