Repository: THUDM/CogDL (language: python)

Issue:
Add GraphSAGE (sample) on Reddit
The current GraphSAGE implementation sample neighbors for **every** nodes in one batch. This only works for toy datasets (e.g., Cora), not for larger datasets (e.g., Reddit)


Discussion / hints:

Fix patch:
diff --git a/cogdl/data/sampler.py b/cogdl/data/sampler.py
index 7f6f1fa8..63b705b5 100644
--- a/cogdl/data/sampler.py
+++ b/cogdl/data/sampler.py
@@ -19,7 +19,11 @@ class Sampler:
     def __init__(self, data, args_params):
         self.data = data.clone()
         self.num_nodes = self.data.x.size()[0]
-        self.num_edges = self.data.edge_index.size()[1]
+        self.num_edges = (
+            self.data.edge_index_train.size()[1]
+            if hasattr(self.data, "edge_index_train")
+            else self.data.edge_index.size()[1]
+        )
 
     def sample(self):
         pass
@@ -41,7 +45,11 @@ class SAINTSampler(Sampler):
 
     def __init__(self, data, args_params):
         super().__init__(data, args_params)
-        edge_index = self.data.edge_index.cpu().numpy()
+        edge_index = (
+            self.data.edge_index_train.cpu().numpy()
+            if hasattr(self.data, "edge_index_train")
+            else self.data.edge_index.cpu().numpy()
+        )
         self.adj = sp.coo_matrix(
             (np.ones(self.num_edges), (edge_index[0], edge_index[1])),
             shape=(self.num_nodes, self.num_nodes),
@@ -113,7 +121,11 @@ def sample(self):
 
     def extract_subgraph(self, edge_idx, directed=True):
         edge_idx = np.unique(edge_idx)
-        subg_edge = self.data.edge_index.t()[edge_idx].cpu().numpy()
+        subg_edge = (
+            self.data.edge_index_train.t()[edge_idx].cpu().numpy()
+            if hasattr(self.data, "edge_index_train")
+            else self.data.edge_index.t()[edge_idx].cpu().numpy()
+        )
         if not directed:
             subg_edge = np.concatenate((subg_edge, subg_edge[:, [1, 0]]))
         subg_edge = np.unique(subg_edge, axis=0)
@@ -168,7 +180,7 @@ def get_subgraph(self, phase, require_norm=True):
             node_subgraph = np.arange(self.data.num_nodes)
             data = self.data.clone()
             if require_norm:
-                data.norm_aggr = torch.ones(self.num_edges)
+                data.norm_aggr = torch.ones(self.data.edge_index.size()[1])
                 data.norm_loss = self.norm_loss_test
         else:
             if len(self.subgraphs_nodes) == 0:
diff --git a/cogdl/tasks/attributed_graph_clustering.py b/cogdl/tasks/attributed_graph_clustering.py
index bfafcf05..cc92a55c 100644
--- a/cogdl/tasks/attributed_graph_clustering.py
+++ b/cogdl/tasks/attributed_graph_clustering.py
@@ -11,6 +11,7 @@
 
 from . import BaseTask, register_task
 
+
 @register_task("attributed_graph_clustering")
 class AttributedGraphClustering(BaseTask):
     """Attributed graph clustring task."""
@@ -24,6 +25,7 @@ def add_args(parser: argparse.ArgumentParser):
         parser.add_argument("--cluster-method", type=str, default="kmeans")
         parser.add_argument("--hidden-size", type=int, default=128)
         parser.add_argument("--model-type", type=str, default="content")
+        parser.add_argument("--evaluate", type=str, default="full")
         # fmt: on
 
     def __init__(
@@ -49,14 +51,16 @@ def __init__(
             self.hidden_size = args.hidden_size = args.hidden_size
             args.num_features = dataset.num_features
         self.model = build_model(args)
-        
         self.num_clusters = args.num_clusters
         if args.cluster_method not in ["kmeans", "spectral"]:
             raise Exception("cluster method must be kmeans or spectral")
         if args.model_type not in ["content", "spectral", "both"]:
             raise Exception("model type must be content, spectral or both")
         self.cluster_method = args.cluster_method
+        if args.evaluate not in ["full", "NMI"]:
+            raise Exception("evaluation must be full or NMI")
         self.model_type = args.model_type
+        self.evaluate = args.evaluate
         self.is_weighted = self.data.edge_attr is not None
         self.enhance = args.enhance
 
@@ -82,12 +86,12 @@ def train(self) -> Dict[str, float]:
                 features_matrix[node] = embeddings[vid]
             features_matrix = torch.tensor(features_matrix)
             features_matrix = F.normalize(features_matrix, p=2, dim=1)
-            #features_matrix = self.momentum * node_attr + (1 - self.momentum) * features_matrix
+            # features_matrix = self.momentum * node_attr + (1 - self.momentum) * features_matrix
         else:
             trainer = self.model.get_trainer(AttributedGraphClustering, self.args)(self.args)
             self.model = trainer.fit(self.model, self.data)
             features_matrix = self.model.get_features(self.data)
-            
+
         features_matrix = features_matrix.cpu().numpy()
         print("Clustering...")
         if self.cluster_method == "kmeans":
@@ -96,32 +100,38 @@ def train(self) -> Dict[str, float]:
         else:
             # features_matrix = np.dot(features_matrix, features_matrix.transpose())
             # features_matrix = 0.5 * (np.abs(features_matrix) + np.abs(features_matrix.transpose()))
-            clustering = SpectralClustering(n_clusters=self.num_clusters, assign_labels="discretize", random_state=0).fit(
-                features_matrix
-            )
+            clustering = SpectralClustering(
+                n_clusters=self.num_clusters, assign_labels="discretize", random_state=0
+            ).fit(features_matrix)
             clusters = clustering.labels_
-        return self.__evaluate(clusters)
+        if self.evaluate == "full":
+            return self.__evaluate(clusters, True)
+        else:
+            return self.__evaluate(clusters, False)
 
-    def __evaluate(self, clusters) -> Dict[str, float]:
+    def __evaluate(self, clusters, full=True) -> Dict[str, float]:
         print("Evaluating...")
-        TP = 0
-        FP = 0
-        TN = 0
-        FN = 0
         truth = self.data.y.cpu().numpy()
-        for i in range(self.num_nodes):
-            for j in range(i + 1, self.num_nodes):
-                if clusters[i] == clusters[j] and truth[i] == truth[j]:
-                    TP += 1
-                if clusters[i] != clusters[j] and truth[i] == truth[j]:
-                    FP += 1
-                if clusters[i] == clusters[j] and truth[i] != truth[j]:
-                    FN += 1
-                if clusters[i] != clusters[j] and truth[i] != truth[j]:
-                    TN += 1
-        _ = (TP + TN) / (TP + FP + TN + FN)
-        precision = TP / (TP + FP)
-        recall = TP / (TP + FN)
-        print("TP", TP, "FP", FP, "TN", TN, "FN", FN)
-        micro_f1 = 2 * (precision * recall) / (precision + recall)
-        return dict(Accuracy=precision, NMI=normalized_mutual_info_score(clusters, truth), Micro_F1=micro_f1)
+        if full:
+            TP = 0
+            FP = 0
+            TN = 0
+            FN = 0
+            for i in range(self.num_nodes):
+                for j in range(i + 1, self.num_nodes):
+                    if clusters[i] == clusters[j] and truth[i] == truth[j]:
+                        TP += 1
+                    if clusters[i] != clusters[j] and truth[i] == truth[j]:
+                        FP += 1
+                    if clusters[i] == clusters[j] and truth[i] != truth[j]:
+                        FN += 1
+                    if clusters[i] != clusters[j] and truth[i] != truth[j]:
+                        TN += 1
+            _ = (TP + TN) / (TP + FP + TN + FN)
+            precision = TP / (TP + FP)
+            recall = TP / (TP + FN)
+            print("TP", TP, "FP", FP, "TN", TN, "FN", FN)
+            micro_f1 = 2 * (precision * recall) / (precision + recall)
+            return dict(Accuracy=precision, NMI=normalized_mutual_info_score(clusters, truth), Micro_F1=micro_f1)
+        else:
+            return dict(NMI=normalized_mutual_info_score(clusters, truth))
diff --git a/cogdl/trainers/agc_trainer.py b/cogdl/trainers/agc_trainer.py
index 8775d512..260025a1 100644
--- a/cogdl/trainers/agc_trainer.py
+++ b/cogdl/trainers/agc_trainer.py
@@ -73,6 +73,8 @@ def compute_intra(self, x, clusters):
         # print(intra / num_per_cluster)
         return np.mean(intra / num_per_cluster)
 
+
+"""
     def evaluate(self, clusters, truth):
         print("Evaluating...")
         TP = 0
@@ -96,3 +98,4 @@ def evaluate(self, clusters, truth):
         print(
             "Accuracy = ", precision, "NMI = ", normalized_mutual_info_score(clusters, truth), "Micro_F1 = ", micro_f1
         )
+"""
diff --git a/cogdl/trainers/daegc_trainer.py b/cogdl/trainers/daegc_trainer.py
index f87ca21a..e5841733 100644
--- a/cogdl/trainers/daegc_trainer.py
+++ b/cogdl/trainers/daegc_trainer.py
@@ -98,27 +98,3 @@ def getP(self, Q):
     def cluster_loss(self, P, Q):
         # return nn.MSELoss(reduce=True, size_average=False)(P, Q)
         return nn.KLDivLoss(reduce=True, size_average=False)(P.log(), Q)
-
-    def evaluate(self, clusters, truth):
-        print("Evaluating...")
-        TP = 0
-        FP = 0
-        TN = 0
-        FN = 0
-        for i in range(self.num_nodes):
-            for j in range(i + 1, self.num_nodes):
-                if clusters[i] == clusters[j] and truth[i] == truth[j]:
-                    TP += 1
-                if clusters[i] != clusters[j] and truth[i] == truth[j]:
-                    FP += 1
-                if clusters[i] == clusters[j] and truth[i] != truth[j]:
-                    FN += 1
-                if clusters[i] != clusters[j] and truth[i] != truth[j]:
-                    TN += 1
-        precision = TP / (TP + FP)
-        recall = TP / (TP + FN)
-        print("TP", TP, "FP", FP, "TN", TN, "FN", FN)
-        micro_f1 = 2 * (precision * recall) / (precision + recall)
-        print(
-            "Accuracy = ", precision, "NMI = ", normalized_mutual_info_score(clusters, truth), "Micro_F1 = ", micro_f1
-        )


Test patch:
diff --git a/tests/tasks/test_attributed_graph_clustering.py b/tests/tasks/test_attributed_graph_clustering.py
index b295958f..3d74de8a 100644
--- a/tests/tasks/test_attributed_graph_clustering.py
+++ b/tests/tasks/test_attributed_graph_clustering.py
@@ -13,6 +13,7 @@ def get_default_args():
         "device_id": [0],
         "num_clusters": 7,
         "cluster_method": "kmeans",
+        "evaluate": "NMI",
         "hidden_size": 16,
         "model_type": "spectral",
         "enhance": None,
@@ -21,12 +22,12 @@ def get_default_args():
         "theta": 0.5,
         "mu": 0.2,
         "checkpoint": False,
-        "walk_length": 80,
-        "walk_num": 40,
+        "walk_length": 10,
+        "walk_num": 4,
         "window_size": 5,
-        "worker": 10,
+        "worker": 2,
         "iteration": 3,
-        "rank": 256,
+        "rank": 64,
         "negative": 1,
         "is_large": False,
         "max_iter": 5,
@@ -61,7 +62,7 @@ def test_kmeans_cora():
     args.cluster_method = "kmeans"
     task = build_task(args)
     ret = task.train()
-    assert ret["Accuracy"] > 0
+    assert ret["NMI"] > 0
 
 
 def test_kmeans_citeseer():
@@ -73,7 +74,7 @@ def test_kmeans_citeseer():
     args.num_clusters = 6
     task = build_task(args)
     ret = task.train()
-    assert ret["Accuracy"] > 0
+    assert ret["NMI"] > 0
 
 
 def test_spectral_cora():
@@ -84,7 +85,7 @@ def test_spectral_cora():
     args.cluster_method = "spectral"
     task = build_task(args)
     ret = task.train()
-    assert ret["Accuracy"] > 0
+    assert ret["NMI"] > 0
 
 
 def test_spectral_citeseer():
@@ -96,7 +97,7 @@ def test_spectral_citeseer():
     args.num_clusters = 6
     task = build_task(args)
     ret = task.train()
-    assert ret["Accuracy"] > 0
+    assert ret["NMI"] > 0
 
 
 def test_prone_cora():
@@ -107,7 +108,7 @@ def test_prone_cora():
     args.cluster_method = "kmeans"
     task = build_task(args)
     ret = task.train()
-    assert ret["Accuracy"] > 0
+    assert ret["NMI"] > 0
 
 
 def test_deepwalk_cora():
@@ -118,7 +119,7 @@ def test_deepwalk_cora():
     args.cluster_method = "kmeans"
     task = build_task(args)
     ret = task.train()
-    assert ret["Accuracy"] > 0
+    assert ret["NMI"] > 0
 
 
 def test_netmf_cora():
@@ -129,7 +130,7 @@ def test_netmf_cora():
     args.cluster_method = "kmeans"
     task = build_task(args)
     ret = task.train()
-    assert ret["Accuracy"] > 0
+    assert ret["NMI"] > 0
 
 
 def test_agc_cora():
@@ -138,10 +139,10 @@ def test_agc_cora():
     args.model_type = "both"
     args.dataset = "cora"
     args.cluster_method = "spectral"
-    args.max_iter = 10
+    args.max_iter = 2
     task = build_task(args)
     ret = task.train()
-    assert ret["Accuracy"] > 0
+    assert ret["NMI"] > 0
 
 
 def test_daegc_cora():
@@ -152,7 +153,7 @@ def test_daegc_cora():
     args.cluster_method = "kmeans"
     task = build_task(args)
     ret = task.train()
-    assert ret["Accuracy"] > 0
+    assert ret["NMI"] > 0
 
 
 if __name__ == "__main__":