feat: add evaluate_from_embeddings

Author: galogm

graph_datasets/utils/evaluation/eval_tools.py

@@ -17,6 +17,7 @@
 from sklearn.metrics import normalized_mutual_info_score as NMI
 from sklearn.svm import LinearSVC
 
+from ..common import get_str_time
 from ..common import tab_printer
 
 
@@ -38,7 +39,14 @@ def save_dict(di_, filename_):
         pickle.dump(di_, f)
 
 
-def split_train_test_nodes(data, train_ratio, valid_ratio, data_name, split_id=0, fixed_split=True):
+def split_train_test_nodes(
+    num_nodes,
+    train_ratio,
+    valid_ratio,
+    data_name,
+    split_id=0,
+    fixed_split=True,
+):
     if fixed_split:
         file_path = f"../input/fixed_splits/{data_name}-{train_ratio}-{valid_ratio}-splits.npy"
         if not os.path.exists(file_path):
@@ -47,13 +55,12 @@ def split_train_test_nodes(data, train_ratio, valid_ratio, data_name, split_id=0
             splits = {}
             for idx in range(10):
                 # set up train val and test
-                shuffle = list(range(data.num_nodes))
+                shuffle = list(range(num_nodes))
                 random.shuffle(shuffle)
-                train_nodes = shuffle[:int(data.num_nodes * train_ratio / 100)]
-                val_nodes = shuffle[
-                    int(data.num_nodes * train_ratio /
-                        100):int(data.num_nodes * (train_ratio + valid_ratio) / 100)]
-                test_nodes = shuffle[int(data.num_nodes * (train_ratio + valid_ratio) / 100):]
+                train_nodes = shuffle[:int(num_nodes * train_ratio / 100)]
+                val_nodes = shuffle[int(num_nodes * train_ratio /
+                                        100):int(num_nodes * (train_ratio + valid_ratio) / 100)]
+                test_nodes = shuffle[int(num_nodes * (train_ratio + valid_ratio) / 100):]
                 splits[idx] = {"train": train_nodes, "valid": val_nodes, "test": test_nodes}
             save_dict(di_=splits, filename_=file_path)
         else:
@@ -62,12 +69,12 @@ def split_train_test_nodes(data, train_ratio, valid_ratio, data_name, split_id=0
         train_nodes, val_nodes, test_nodes = split["train"], split["valid"], split["test"]
     else:
         # set up train val and test
-        shuffle = list(range(data.num_nodes))
+        shuffle = list(range(num_nodes))
         random.shuffle(shuffle)
-        train_nodes = shuffle[:int(data.num_nodes * train_ratio / 100)]
-        val_nodes = shuffle[int(data.num_nodes * train_ratio /
-                                100):int(data.num_nodes * (train_ratio + valid_ratio) / 100)]
-        test_nodes = shuffle[int(data.num_nodes * (train_ratio + valid_ratio) / 100):]
+        train_nodes = shuffle[:int(num_nodes * train_ratio / 100)]
+        val_nodes = shuffle[int(num_nodes * train_ratio /
+                                100):int(num_nodes * (train_ratio + valid_ratio) / 100)]
+        test_nodes = shuffle[int(num_nodes * (train_ratio + valid_ratio) / 100):]
 
     return np.array(train_nodes), np.array(val_nodes), np.array(test_nodes)
 
@@ -175,18 +182,18 @@ def kmeans_test(X, y, n_clusters, repeat=10):
     )
 
 
-def svm_test(data, embeddings, labels, train_ratios=(10, 20, 30, 40), repeat=10):
+def svm_test(num_nodes, data_name, embeddings, labels, train_ratios=(10, 20, 30, 40), repeat=10):
     result_macro_f1_list = []
     result_micro_f1_list = []
     for train_ratio in train_ratios:
         macro_f1_list = []
         micro_f1_list = []
         for i in range(repeat):
             train_idx, val_idx, test_idx = split_train_test_nodes(
-                data=data,
+                num_nodes=num_nodes,
                 train_ratio=train_ratio,
                 valid_ratio=train_ratio,
-                data_name=data.name,
+                data_name=data_name,
                 split_id=i,
             )
             X_train, X_test = embeddings[np.concatenate([train_idx, val_idx])], embeddings[test_idx]
@@ -204,16 +211,16 @@ def svm_test(data, embeddings, labels, train_ratios=(10, 20, 30, 40), repeat=10)
 
 
 def evaluate_results_nc(
-    data,
+    labels,
+    num_classes,
+    num_nodes,
+    data_name,
     embeddings,
     quiet=False,
     method="unsup",
     alpha: float = 2.0,
     beta: float = 2.0,
 ):
-    labels = data.y.detach().cpu().numpy()
-    num_classes = data.num_classes
-    num_nodes = data.num_nodes
     if embeddings.shape[0] > num_nodes:
         z_1 = embeddings[:num_nodes]
         z_2 = embeddings[num_nodes:]
@@ -227,7 +234,8 @@ def evaluate_results_nc(
             svm_macro_f1_list,
             svm_micro_f1_list,
         ) = svm_test(
-            data=data,
+            num_nodes=num_nodes,
+            data_name=data_name,
             embeddings=embeddings,
             labels=labels,
         )
@@ -315,7 +323,7 @@ def save_embedding(
     verbose: bool or int = True,
 ):
     dataset_name = dataset_name.replace("_", "-")
-    timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
+    timestamp = get_str_time()
     file_name = f"{dataset_name.lower()}_{model_name.lower()}_embeds_{timestamp}.pth"
     file_path = os.path.join(save_dir, file_name)
 

graph_datasets/utils/evaluation/evaluation.py

@@ -5,6 +5,7 @@
 from typing import Dict
 from typing import Tuple
 
+import numpy as np
 import torch
 
 from .eval_tools import evaluate_results_nc
@@ -22,44 +23,31 @@ def load_from_file(file_name):
     return embeddings
 
 
-def evaluate_from_embed_file(
-    embedding_file: str,
-    data_file: str,
-    save_path: str = "./tmp/",
+def evaluate_from_embeddings(
+    labels: np.ndarray,
+    num_classes: int,
+    num_nodes: int,
+    data_name: str,
+    embeddings: torch.Tensor,
     quiet: bool = True,
+    method: str = "both",
 ) -> Tuple[Dict, Dict]:
-    """Evaluation of representation quality using clustering and classification tasks.
+    """evaluate embeddings with LR and Clustering.
 
     Args:
-        embedding_file (str): Embedded file name.
-        data_file (str): Data file name.
-        save_path (str, optional): Folder path to store. Defaults to './tmp/'.
-        quiet (bool, optional): Whether to print results. Defaults to True.
+        labels (np.ndarray): labels.
+        num_classes (int): number of classes.
+        num_nodes (int): number of nodes.
+        data_name (str): name of the datasets.
+        embeddings (torch.Tensor): embeddings.
+        quiet (bool, optional): whether to print info. Defaults to True.
+        method (bool, optional): method for evaluation, \
+            "sup" for linear regression, "unsup" for svm clustering, "both" for both.\
+            Defaults to "both".
 
     Returns:
-        Tuple[Dict, Dict]: Two dicts are included, \
-            which are the evaluation results of clustering and classification.
-
-    Example:
-        .. code-block:: python
-
-            from graph_datasets import evaluate_from_embed_file
-
-            method_name='orderedgnn'
-            data_name='texas'
-
-            clustering_res, classification_res = evaluate_from_embed_file(
-                f'{data_name}_{method_name}_embeds.pth',
-                f'{data_name}_data.pth',
-                save_path='./save/',
-            )
+        Tuple[Dict, Dict]: (clustering_results, classification_results)
     """
-    embedding_file = os.path.join(save_path, embedding_file)
-    data_file = os.path.join(save_path, data_file)
-
-    embeddings = load_from_file(embedding_file).cpu().detach()
-    data = load_from_file(data_file)
-
     # Call the evaluate_results_nc function with the loaded embeddings
     (
         svm_macro_f1_list,
@@ -75,10 +63,13 @@ def evaluate_from_embed_file(
         f1_mean,
         f1_std,
     ) = evaluate_results_nc(
-        data,
+        labels,
+        num_classes,
+        num_nodes,
+        data_name,
         embeddings,
         quiet=quiet,
-        method="both",
+        method=method,
     )
 
     # Format the output as desired
@@ -87,20 +78,73 @@ def evaluate_from_embed_file(
         "NMI": f"{nmi_mean * 100:.2f}±{nmi_std * 100:.2f}",
         "AMI": f"{ami_mean * 100:.2f}±{ami_std * 100:.2f}",
         "ARI": f"{ari_mean * 100:.2f}±{ari_std * 100:.2f}",
-        "Macro F1": f"{f1_mean * 100:.2f}±{f1_std * 100:.2f}",
+        "MaF1": f"{f1_mean * 100:.2f}±{f1_std * 100:.2f}",
     }
 
     svm_macro_f1_list = [f"{res[0] * 100:.2f}±{res[1] * 100:.2f}" for res in svm_macro_f1_list]
     svm_micro_f1_list = [f"{res[0] * 100:.2f}±{res[1] * 100:.2f}" for res in svm_micro_f1_list]
 
     classification_results = {}
     for i, percent in enumerate(["10%", "20%", "30%", "40%"]):
-        classification_results[f"{percent}_Macro-F1"] = svm_macro_f1_list[i]
-        classification_results[f"{percent}_Micro-F1"] = svm_micro_f1_list[i]
+        classification_results[f"{percent}_MaF1"] = svm_macro_f1_list[i]
+        classification_results[f"{percent}_MiF1"] = svm_micro_f1_list[i]
 
     return clustering_results, classification_results
 
 
+def evaluate_from_embed_file(
+    embedding_file: str,
+    data_file: str,
+    save_path: str = "./tmp/",
+    quiet: bool = True,
+) -> Tuple[Dict, Dict]:
+    """Evaluation of representation quality using clustering and classification tasks.
+
+    Args:
+        embedding_file (str): Embedded file name.
+        data_file (str): Data file name.
+        save_path (str, optional): Folder path to store. Defaults to './tmp/'.
+        quiet (bool, optional): Whether to print results. Defaults to True.
+
+    Returns:
+        Tuple[Dict, Dict]: Two dicts are included, \
+            which are the evaluation results of clustering and classification.
+
+    Example:
+        .. code-block:: python
+
+            from graph_datasets import evaluate_from_embed_file
+
+            method_name='orderedgnn'
+            data_name='texas'
+
+            clustering_res, classification_res = evaluate_from_embed_file(
+                f'{data_name}_{method_name}_embeds.pth',
+                f'{data_name}_data.pth',
+                save_path='./save/',
+            )
+    """
+    embedding_file = os.path.join(save_path, embedding_file)
+    data_file = os.path.join(save_path, data_file)
+
+    embeddings = load_from_file(embedding_file).cpu().detach()
+    data = load_from_file(data_file)
+
+    labels = data.y.detach().cpu().numpy()
+    num_classes = data.num_classes
+    num_nodes = data.num_nodes
+    data_name = data.name
+
+    return evaluate_from_embeddings(
+        labels=labels,
+        num_classes=num_classes,
+        num_nodes=num_nodes,
+        data_name=data_name,
+        embeddings=embeddings,
+        quiet=quiet,
+    )
+
+
 # if __name__ == "__main__":
 #     method_name = 'orderedgnn'  # 'selene' 'greet' 'hgrl' 'nwr-gae' 'orderedgnn'
 #     data_name = 'texas'  # 'actor' 'chameleon' 'cornell' 'squirrel' 'texas' 'wisconsin'