fix edge index in run_model

Author: mishadr

experiments/various_tasks.py

@@ -108,7 +108,7 @@ def link_prediction():
 
     gen_dataset = DatasetManager.get_by_config(dc, dvc)
     print(gen_dataset.data)
-    gen_dataset.train_test_split(percent_train_class=0.85, percent_test_class=0.1)
+    gen_dataset.train_test_split(percent_train_class=0.8, percent_test_class=0.2)
 
     gnn = FrameworkGNNConstructor(
         model_config=ModelConfig(
@@ -227,7 +227,8 @@ def link_prediction():
     edge_label_index = torch.tensor([[5], [6]])
 
     # get embeddings for all nodes
-    node_out = gnn(data.x, data.edge_index)
+    train_edge_index = gen_dataset.edge_label_index[:, gen_dataset.train_mask]
+    node_out = gnn(data.x, train_edge_index)
 
     # Get embeddings for our nodes
     src = node_out[edge_label_index[0]]

gnn_aid/datasets/gen_dataset.py

@@ -366,7 +366,8 @@ def train_test_split(
                      train_data.edge_label_index.size(1) + val_data.edge_label_index.size(1)] = True
 
             test_mask = torch.zeros(total_edges, dtype=torch.bool)
-            test_mask[-test_data.edge_label_index.size(1):] = True
+            if test_data.edge_label_index.size(1) > 0:
+                test_mask[-test_data.edge_label_index.size(1):] = True
         else:
             raise ValueError(f"Unsupported task type {task_type}")
 

gnn_aid/models_builder/model_managers/framework_mm.py

@@ -519,13 +519,12 @@ def run_model(
 
             elif task_type.is_edge_level():
                 data = gen_dataset.data
-                edge_label_index = mask_tensor
-
+                train_edge_index = gen_dataset.edge_label_index[:, gen_dataset.train_mask]
                 data_x_copy = torch.clone(data.x)
 
                 # FIXME misha check, test
                 if hasattr(self, 'mask_features'):
-                    node_ind = torch.unique(edge_label_index)
+                    node_ind = torch.unique(train_edge_index)
                     for elem_ind in node_ind:
                         for feature in self.mask_features:
                             data_x_copy[elem_ind][gen_dataset.node_attr_slices[feature][0]:
@@ -535,21 +534,23 @@ def run_model(
                 if hasattr(self, 'optimizer'):
                     self.optimizer.zero_grad()
 
-                # get logits for nodes
-                node_out = self.gnn(data_x_copy, data.edge_index)
-
-                src = node_out[edge_label_index[0]]
-                dst = node_out[edge_label_index[1]]
+                # get logits for all nodes based on train edges
+                node_out = self.gnn(data_x_copy, train_edge_index)
 
+                # Get logits for edges from mask
+                src = node_out[mask_tensor[0]]
+                dst = node_out[mask_tensor[1]]
                 edge_out = self.gnn.decode(src, dst)
 
+                # Apply different out
                 full_out = None
                 if out == 'logits':
                     full_out = edge_out
                 elif out == 'predictions':
                     if task_type == Task.EDGE_PREDICTION:
-                        # TODO misha
-                        raise NotImplementedError
+                        # TODO misha is it ok?
+                        full_out = edge_out.softmax(dim=-1)
+                        # raise NotImplementedError
                     elif task_type == Task.EDGE_CLASSIFICATION:
                         full_out = edge_out.softmax(dim=-1)
                     elif task_type == Task.EDGE_REGRESSION:

gnn_aid/models_builder/models_utils.py

@@ -255,6 +255,11 @@ def name(
             res += '{' + kwargs + '}'
         return res
 
+    def __str__(
+            self
+    ) -> str:
+        return self.name()
+
     def compute(
             self,
             y_true,