turn back functions in GCN.py

Author: Demoren1

code/GCN.py

@@ -241,8 +241,7 @@ def train_model_diversity(
     model,
     train_loader,
     valid_loader,
-    discrete_diversity_matrix_train,
-    discrete_diversity_matrix_valid,
+    discrete_diversity_matrix,
     optimizer,
     criterion,
     num_epochs,
@@ -271,7 +270,7 @@ def train_model_diversity(
             anchor = model(features, edge_index)
 
             positive_index, negative_index = get_positive_and_negative(
-                discrete_diversity_matrix_train, i
+                discrete_diversity_matrix, index, train_loader.dataset
             )
             if positive_index is None or negative_index is None:
                 continue
@@ -318,7 +317,7 @@ def train_model_diversity(
                 anchor = model(features, edge_index)
 
                 positive_index, negative_index = get_positive_and_negative(
-                    discrete_diversity_matrix_valid, i
+                    discrete_diversity_matrix, index, valid_loader.dataset
                 )
                 if positive_index is None or negative_index is None:
                     continue
@@ -478,8 +477,28 @@ def train_model_accuracy(
     return train_losses, valid_losses
 
 
-def get_positive_and_negative(diversity_matrix, index):
-    positive = np.where(diversity_matrix[index] == 1)[0]
-    negative = np.where(diversity_matrix[index] == 0)[0]
+def get_positive_and_negative(diversity_matrix, index, dataset=None):
+    positive = np.where(
+        (diversity_matrix[index, :] == 1) & (np.arange(len(diversity_matrix)) != index)
+    )[0].tolist()
+    negative = np.where(diversity_matrix[index, :] == -1)[0].tolist()
+
+    if dataset is not None:
+        appropriate_indexes = [dataset[i][2] for i in range(len(dataset))]
+
+        positive = [
+            appropriate_indexes.index(idx)
+            for idx in positive
+            if idx in appropriate_indexes
+        ]
+        negative = [
+            appropriate_indexes.index(idx)
+            for idx in negative
+            if idx in appropriate_indexes
+        ]
+
+    if not positive or not negative:
+        print("Both positive and negative samples are empty!")
+        return None, None
 
     return np.random.choice(positive), np.random.choice(negative)

code/create_GCN_dataset.py

@@ -53,11 +53,9 @@ def apply_argmax_to_predictions(data):
 save_dicts_as_json(first_arch_dicts, output_dir)
 
 apply_argmax_to_predictions(first_arch_dicts)
-output_dir = "tmp_dataset"
-save_dicts_as_json(first_arch_dicts, output_dir)
 
 second_arch_dicts = load_json_from_directory("second_dataset")
+second_arch_dicts.extend(first_arch_dicts)
 
-first_arch_dicts.extend(second_arch_dicts)
 output_dir = "third_dataset"
-save_dicts_as_json(first_arch_dicts, output_dir)
+save_dicts_as_json(second_arch_dicts, output_dir)