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/README.md b/README.md
index 3d2030d4..e2c74243 100644
--- a/README.md
+++ b/README.md
@@ -169,11 +169,27 @@ git clone https://github.com/THUDM/cogdl /cogdl
 ```
 
 
+### Attributed graph clustering
+
+For attributed graph clustering, we adopt accuracy, normalized mutual information(NMI) and Micro F1 score. We only evaluate these measures on Cora dataset because we are still working on experiments with Citeseer and Pubmed dataset.
+
+| Rank | Method                                                       |   Accuracy   |   NMI   |   Micro F1   |
+| :--- | :----------------------------------------------------------- | :-------: | :--------: | :--------: |
+| 1 | DAEGC [(Wang et al, IJCAI'19)](https://arxiv.org/pdf/1906.06532.pdf) | 43.59 | **45.96** | **47.70** |
+| 2   | AGC [(Zhang et al, IJCAI'20)](https://arxiv.org/pdf/1906.01210.pdf) | **47.71** | 42.15  | 45.93  |
+| 3    | NetMF [(Qiu et al, WSDM'18)](http://arxiv.org/abs/1710.02971) | 45.48 | 40.92 | 40.20 |
+| 4    | DeepWalk [(Perozzi et al, KDD'14)](http://arxiv.org/abs/1403.6652) | 38.34 | 38.90 | 39.07 |
+| 5    | ProNE [(Zhang et al, IJCAI'19)](https://www.ijcai.org/Proceedings/2019/0594.pdf)  | 31.12 | 30.76  | 33.02  |
+| 6    | LINE [(Tang et al, WWW'15)](http://arxiv.org/abs/1503.03578)  | 26.17 | 25.48  | 28.15  |
+| 7    | KMeans  | 33.47 | 14.55  | 27.11  |
+| 8    | Spectral  | 24.14 | 14.39  | 25.53  |
+
+
 ## Contributing
 
 ### Add Your Own Dataset
 
-If you have a unique and interesting dataset and are willing to publish it, you can submit your dataset via opening an issue in our repository, we will run all suitable methods on your dataset and update our leaderboard. 
+If you have a unique and interesting dataset and are willing to publish it, you can submit your dataset via [opening an issue](https://github.com/THUDM/cogdl/issues) in our repository, we will run all suitable methods on your dataset and update our leaderboard. 
 
 ### Implement Your Own Model
 
diff --git a/checkpoints/pyg_gcn_cora.pt b/checkpoints/pyg_gcn_cora.pt
new file mode 100644
index 00000000..952e18b8
Binary files /dev/null and b/checkpoints/pyg_gcn_cora.pt differ
diff --git a/cogdl/datasets/matlab_matrix.py b/cogdl/datasets/matlab_matrix.py
index 28e38e67..efaf586c 100644
--- a/cogdl/datasets/matlab_matrix.py
+++ b/cogdl/datasets/matlab_matrix.py
@@ -75,6 +75,7 @@ def __init__(self, args=None):
         super(BlogcatalogDataset, self).__init__(path, filename, url)
 
 
+"""
 @register_dataset("flickr")
 class FlickrDataset(MatlabMatrix):
     def __init__(self, args=None):
@@ -82,6 +83,7 @@ def __init__(self, args=None):
         url = "http://leitang.net/code/social-dimension/data/"
         path = osp.join(osp.dirname(osp.realpath(__file__)), "../..", "data", dataset)
         super(FlickrDataset, self).__init__(path, filename, url)
+"""
 
 
 @register_dataset("wikipedia")
diff --git a/cogdl/datasets/pyg_ogb.py b/cogdl/datasets/pyg_ogb.py
index 549c51f6..a7c77624 100644
--- a/cogdl/datasets/pyg_ogb.py
+++ b/cogdl/datasets/pyg_ogb.py
@@ -8,6 +8,7 @@
 from ogb.graphproppred import PygGraphPropPredDataset
 
 from . import register_dataset
+from cogdl.utils import accuracy_evaluator
 
 
 class OGBNDataset(PygNodePropPredDataset):
@@ -85,6 +86,9 @@ def __init__(self, args=None):
             PygNodePropPredDataset(dataset, path)
         super(OGBPapers100MDataset, self).__init__(path, dataset)
 
+    def get_evaluator(self):
+        return accuracy_evaluator()
+
 
 class OGBGDataset(PygGraphPropPredDataset):
     def __init__(self, root, name):
diff --git a/cogdl/datasets/saint_data.py b/cogdl/datasets/saint_data.py
index 1baec021..d7ec711a 100644
--- a/cogdl/datasets/saint_data.py
+++ b/cogdl/datasets/saint_data.py
@@ -6,7 +6,7 @@
 from sklearn.preprocessing import StandardScaler
 
 from cogdl.data import Data, Dataset
-from cogdl.utils import multilabel_evaluator, download_url
+from cogdl.utils import multilabel_evaluator, download_url, multiclass_evaluator
 from .planetoid_data import index_to_mask
 from . import register_dataset
 
@@ -26,10 +26,16 @@ def read_saint_data(folder):
 
     feats = torch.from_numpy(feats).float()
     item = class_map["0"]
-    label_matrix = np.zeros((feats.shape[0], len(item)), dtype=float)
-    for key, val in class_map.items():
-        label_matrix[int(key)] = np.array(val)
-    label_matrix = torch.from_numpy(label_matrix)
+    if isinstance(item, list):
+        labels = np.zeros((feats.shape[0], len(item)), dtype=float)
+        for key, val in class_map.items():
+            labels[int(key)] = np.array(val)
+    else:
+        labels = np.zeros(feats.shape[0], dtype=np.long)
+        for key, val in class_map.items():
+            labels[int(key)] = val
+
+    labels = torch.from_numpy(labels)
 
     def get_adj(adj):
         row, col = adj.nonzero()
@@ -45,7 +51,7 @@ def get_adj(adj):
 
     data = Data(
         x=feats,
-        y=label_matrix,
+        y=labels,
         edge_index=edge_index_full,
         edge_attr=edge_attr_full,
         edge_index_train=edge_index_train,
@@ -76,6 +82,8 @@ def processed_file_names(self):
     @property
     def num_classes(self):
         assert hasattr(self.data, "y")
+        if len(self.data.y.shape) == 1:
+            return int(torch.max(self.data.y) + 1)
         return self.data.y.shape[1]
 
     def download(self):
@@ -103,6 +111,7 @@ def scale_feats(data):
     scaler = StandardScaler()
     scaler.fit(data.x.numpy())
     data.x = torch.from_numpy(scaler.transform(data.x)).float()
+    print(data.x.shape)
     return data
 
 
@@ -128,3 +137,18 @@ def __init__(self, args=None):
             SAINTDataset(path, dataset, url)
         super(AmazonDataset, self).__init__(path, dataset, url)
         self.data = scale_feats(self.data)
+
+
+@register_dataset("flickr")
+class FlickrDatset(SAINTDataset):
+    def __init__(self, args=None):
+        dataset = "Flickr"
+        url = "https://cloud.tsinghua.edu.cn/d/d3ebcb5fa2da463b8213/files/?p=%2F{}&dl=1"
+        path = osp.join(osp.dirname(osp.realpath(__file__)), "../..", "data", dataset)
+        if not osp.exists(path):
+            SAINTDataset(path, dataset, url)
+        super(FlickrDatset, self).__init__(path, dataset, url)
+        self.data = scale_feats(self.data)
+
+    def get_evaluator(self):
+        return multiclass_evaluator()
diff --git a/cogdl/match.yml b/cogdl/match.yml
index c03eb28f..70b336e2 100644
--- a/cogdl/match.yml
+++ b/cogdl/match.yml
@@ -32,8 +32,12 @@ node_classification:
     - citeseer
     - pubmed
     - ogbn-arxiv
+    - ogbn-products
+    - ogbn-papers100M
+    - flickr
+    - amazon
+    - yelp
     - reddit
-    - jknet_cora
 unsupervised_node_classification:
   - model:
     - prone
@@ -185,7 +189,10 @@ attributed_graph_clustering:
     - prone
     - netmf
     - line
+    - agc
+    - daegc
     dataset:
     - cora
     - citeseer
     - pubmed
+    - test_small
diff --git a/cogdl/models/agc/__init__.py b/cogdl/models/agc/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/cogdl/models/agc/agc.py b/cogdl/models/agc/agc.py
new file mode 100644
index 00000000..6e00a43d
--- /dev/null
+++ b/cogdl/models/agc/agc.py
@@ -0,0 +1,38 @@
+import torch
+import torch.nn
+import torch.nn.functional as F
+
+from .. import BaseModel, register_model
+from cogdl.trainers.agc_trainer import AGCTrainer
+
+@register_model("agc")
+class AGC(BaseModel):
+    r"""The AGC model from the `"Attributed Graph Clustering via Adaptive Graph Convolution"
+    <https://arxiv.org/abs/1906.01210>`_ paper
+
+    Args:
+        num_clusters (int) : Number of clusters.
+        max_iter     (int) : Max iteration to increase k
+    """
+    @staticmethod
+    def add_args(parser):
+        parser.add_argument("--max-iter", type=int, default=60)
+
+    @classmethod
+    def build_model_from_args(cls, args):
+        return cls(args.num_clusters, args.max_iter)
+
+    def __init__(self, num_clusters, max_iter):
+        super(AGC, self).__init__()
+
+        self.num_clusters = num_clusters
+        self.max_iter = max_iter
+        self.k = 0
+        self.features_matrix = None
+
+    def get_trainer(self, task, args):
+        return AGCTrainer
+
+    def get_features(self, data):
+        return self.features_matrix.detach().cpu()
+    
\ No newline at end of file
diff --git a/cogdl/models/agc/daegc.py b/cogdl/models/agc/daegc.py
new file mode 100644
index 00000000..e2a8ce92
--- /dev/null
+++ b/cogdl/models/agc/daegc.py
@@ -0,0 +1,90 @@
+import networkx as nx
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch_geometric.nn.conv import GATConv
+
+from .. import BaseModel, register_model
+from cogdl.trainers.daegc_trainer import DAEGCTrainer
+
+@register_model("daegc")
+class DAEGC(BaseModel):
+    r"""The DAEGC model from the `"Attributed Graph Clustering: A Deep Attentional Embedding Approach"
+    <https://arxiv.org/abs/1906.06532>`_ paper
+
+    Args:
+        num_clusters (int) : Number of clusters.
+        T (int) : Number of iterations to recalculate P and Q
+        gamma (float) : Hyperparameter that controls two parts of the loss.
+    """
+    @staticmethod
+    def add_args(parser):
+        """Add model-specific arguments to the parser."""
+        # fmt: off
+        parser.add_argument("--num-features", type=int)
+        parser.add_argument("--hidden-size", type=int, default=256)
+        parser.add_argument("--embedding-size", type=int, default=16)
+        parser.add_argument("--num-heads", type=int, default=1)
+        parser.add_argument("--dropout", type=float, default=0)
+        parser.add_argument("--max-epoch", type=int, default=100)
+        parser.add_argument("--lr", type=float, default=0.001)
+        parser.add_argument("--T", type=int, default=5)
+        parser.add_argument("--gamma", type=float, default=10)
+        # fmt: on
+
+    @classmethod
+    def build_model_from_args(cls, args):
+        return cls(
+            args.num_features,
+            args.hidden_size, 
+            args.embedding_size,
+            args.num_heads,
+            args.dropout,
+            args.num_clusters
+        )
+
+    def __init__(self, num_features, hidden_size, embedding_size, num_heads, dropout, num_clusters):
+        super(DAEGC, self).__init__()
+
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        self.embedding_size = embedding_size
+        self.dropout = dropout
+        self.num_clusters = num_clusters
+        self.conv1 = GATConv(num_features, hidden_size, heads=num_heads, dropout=dropout)
+        self.conv2 = GATConv(hidden_size * num_heads, embedding_size, dropout=dropout)
+        self.cluster_center = None
+
+    def get_trainer(self, task, args):
+        return DAEGCTrainer
+
+    def forward(self, x, edge_index):
+        x = F.dropout(x, p=self.dropout, training=self.training)
+        x = F.relu(self.conv1(x, edge_index))
+        x = F.dropout(x, p=self.dropout, training=self.training)
+        x = F.elu(self.conv2(x, edge_index))
+        return F.normalize(x, p=2, dim=1)
+
+    def get_2hop(self, edge_index):
+        r"""add 2-hop neighbors as new edges"""
+        G = nx.Graph()
+        G.add_edges_from(edge_index.t().tolist())
+        H = nx.Graph()
+        for i in range(G.number_of_nodes()):
+            layers = dict(nx.bfs_successors(G, source=i, depth_limit=2))
+            for succ in layers:
+                for idx in layers[succ]:
+                    H.add_edge(i, idx)
+        return torch.tensor(list(H.edges())).t()
+
+    def get_features(self, data):
+        return self.forward(data.x, data.edge_index).detach()
+
+    def recon_loss(self, z, adj):
+        #print(torch.mm(z, z.t()), adj)
+        return F.binary_cross_entropy(
+            F.softmax(torch.mm(z, z.t())), 
+            adj,
+            reduction='sum'
+        )
diff --git a/cogdl/models/nn/pyg_gcn.py b/cogdl/models/nn/pyg_gcn.py
index ed177e07..3b7c29de 100644
--- a/cogdl/models/nn/pyg_gcn.py
+++ b/cogdl/models/nn/pyg_gcn.py
@@ -5,6 +5,7 @@
 
 from .. import BaseModel, register_model
 from cogdl.trainers.sampled_trainer import SAINTTrainer
+from cogdl.trainers.self_task_trainer import SelfTaskTrainer
 
 
 @register_model("pyg_gcn")
@@ -24,6 +25,9 @@ def add_args(parser):
         parser.add_argument('--num-walks', default=50, type=int, help='number of random walks')
         parser.add_argument('--walk-length', default=20, type=int, help='random walk length')
         parser.add_argument('--size-frontier', default=20, type=int, help='frontier size in multidimensional random walks')
+        parser.add_argument('--auxiliary-task', default="none", type=str)
+        parser.add_argument('--alpha', default=10, type=float)
+        parser.add_argument('--label-mask', default=0, type=float)
         # fmt: on
 
     @classmethod
@@ -39,6 +43,8 @@ def build_model_from_args(cls, args):
     def get_trainer(self, task, args):
         if args.sampler != "none":
             return SAINTTrainer
+        elif args.auxiliary_task != 'none':
+            return SelfTaskTrainer
         else:
             return None
 
@@ -62,6 +68,12 @@ def forward(self, x, edge_index, weight=None):
         x = self.convs[-1](x, edge_index, weight)
         return F.log_softmax(x, dim=1)
 
+    def get_embeddings(self, x, edge_index, weight = None):
+        for conv in self.convs[:-1]:
+            x = F.relu(conv(x, edge_index, weight))
+            x = F.dropout(x, p=self.dropout, training=self.training)
+        return x
+
     def node_classification_loss(self, data):
         return F.nll_loss(
             self.forward(data.x, data.edge_index, None if "norm_aggr" not in data else data.norm_aggr)[data.train_mask],
diff --git a/cogdl/tasks/attributed_graph_clustering.py b/cogdl/tasks/attributed_graph_clustering.py
index 80fcec15..bfafcf05 100644
--- a/cogdl/tasks/attributed_graph_clustering.py
+++ b/cogdl/tasks/attributed_graph_clustering.py
@@ -6,13 +6,11 @@
 from sklearn.metrics.cluster import normalized_mutual_info_score
 import torch
 import torch.nn.functional as F
-
 from cogdl.datasets import build_dataset
 from cogdl.models import build_model
 
 from . import BaseTask, register_task
 
-
 @register_task("attributed_graph_clustering")
 class AttributedGraphClustering(BaseTask):
     """Attributed graph clustring task."""
@@ -25,8 +23,7 @@ def add_args(parser: argparse.ArgumentParser):
         parser.add_argument("--num-clusters", type=int, default=7)
         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="emb")
-        parser.add_argument("--momentum", type=float, default=0)
+        parser.add_argument("--model-type", type=str, default="content")
         # fmt: on
 
     def __init__(
@@ -46,22 +43,27 @@ def __init__(
         self.data = dataset[0]
         self.num_nodes = self.data.y.shape[0]
 
-        self.hidden_size = args.hidden_size = dataset.num_features
+        if args.model == "prone":
+            self.hidden_size = args.hidden_size = args.num_features = 13
+        else:
+            self.hidden_size = args.hidden_size = args.hidden_size
+            args.num_features = dataset.num_features
         self.model = build_model(args)
-
-        self.momentum = args.momentum
+        
         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 ["emb", "nn"]:
-            raise Exception("model type must be emb or nn")
+        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
         self.model_type = args.model_type
         self.is_weighted = self.data.edge_attr is not None
         self.enhance = args.enhance
 
     def train(self) -> Dict[str, float]:
-        if self.model_type == "emb":
+        if self.model_type == "content":
+            features_matrix = self.data.x
+        elif self.model_type == "spectral":
             G = nx.Graph()
             if self.is_weighted:
                 edges, weight = (
@@ -80,11 +82,13 @@ 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)
-            node_attr = F.normalize(self.data.x, p=2, dim=1)
-            features_matrix = self.momentum * node_attr + (1 - self.momentum) * features_matrix
-        # TODO: add gnn training methods
-
-        features_matrix = features_matrix.numpy()
+            #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":
             kmeans = KMeans(n_clusters=self.num_clusters, random_state=0).fit(features_matrix)
@@ -92,7 +96,7 @@ 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="kmeans", random_state=0).fit(
+            clustering = SpectralClustering(n_clusters=self.num_clusters, assign_labels="discretize", random_state=0).fit(
                 features_matrix
             )
             clusters = clustering.labels_
@@ -104,7 +108,7 @@ def __evaluate(self, clusters) -> Dict[str, float]:
         FP = 0
         TN = 0
         FN = 0
-        truth = self.data.y.numpy()
+        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]:
diff --git a/cogdl/trainers/agc_trainer.py b/cogdl/trainers/agc_trainer.py
new file mode 100644
index 00000000..8775d512
--- /dev/null
+++ b/cogdl/trainers/agc_trainer.py
@@ -0,0 +1,98 @@
+import torch
+import torch.nn
+import torch.nn.functional as F
+import torch.sparse
+import numpy as np
+from sklearn.cluster import SpectralClustering
+from sklearn.metrics.cluster import normalized_mutual_info_score
+
+from cogdl.utils import add_remaining_self_loops, spmm
+from .base_trainer import BaseTrainer
+
+
+class AGCTrainer(BaseTrainer):
+    def __init__(self, args):
+        self.num_clusters = args.num_clusters
+        self.max_iter = args.max_iter
+        self.device = args.device_id[0] if not args.cpu else "cpu"
+
+    @staticmethod
+    def build_trainer_from_args(args):
+        pass
+
+    def fit(self, model, data):
+        model = model.to(self.device)
+        data.apply(lambda x: x.to(self.device))
+        self.num_nodes = data.x.shape[0]
+
+        adj = data.edge_index
+        adj_values = torch.ones(adj.shape[1]).to(self.device)
+        deg = spmm(adj, adj_values, torch.ones(data.x.shape[0], 1).to(self.device)).squeeze()
+        deg_sqrt = deg.pow(-1 / 2)
+        adj_values = deg_sqrt[adj[1]] * adj_values * deg_sqrt[adj[0]]
+        adj, adj_values = add_remaining_self_loops(adj, adj_values, 1, data.x.shape[0])
+        adj_values = 0.5 * adj_values
+        adj = torch.sparse_coo_tensor(adj, adj_values, torch.Size([data.x.shape[0], data.x.shape[0]])).to(self.device)
+
+        pre_intra = 1e27
+        pre_feat = None
+        for t in range(1, self.max_iter + 1):
+            x = data.x
+            for i in range(t):
+                x = torch.spmm(adj, x)
+            k = torch.mm(x, x.t())
+            w = (torch.abs(k) + torch.abs(k.t())) / 2
+            clustering = SpectralClustering(
+                n_clusters=self.num_clusters, assign_labels="discretize", random_state=0
+            ).fit(w.detach().cpu())
+            clusters = clustering.labels_
+            intra = self.compute_intra(x.cpu().numpy(), clusters)
+            print("iter #%d, intra = %.4lf" % (t, intra))
+            # self.evaluate(clusters, data.y.cpu().numpy())
+            if intra > pre_intra:
+                model.features_matrix = pre_feat
+                model.k = t - 1
+                return model.cpu()
+            pre_intra = intra
+            pre_feat = w
+        model.features_matrix = w
+        model.k = t - 1
+        return model.cpu()
+
+    def compute_intra(self, x, clusters):
+        num_nodes = x.shape[0]
+        intra = np.zeros(self.num_clusters)
+        num_per_cluster = np.zeros(self.num_clusters)
+        for i in range(num_nodes):
+            for j in range(i + 1, num_nodes):
+                if clusters[i] == clusters[j]:
+                    intra[clusters[i]] += np.sum((x[i] - x[j]) ** 2) ** 0.5
+                    num_per_cluster[clusters[i]] += 1
+        intra = np.array(list(filter(lambda x: x > 0, intra)))
+        num_per_cluster = np.array(list(filter(lambda x: x > 0, num_per_cluster)))
+        # print(intra / num_per_cluster)
+        return np.mean(intra / num_per_cluster)
+
+    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
+        )
diff --git a/cogdl/trainers/daegc_trainer.py b/cogdl/trainers/daegc_trainer.py
new file mode 100644
index 00000000..f87ca21a
--- /dev/null
+++ b/cogdl/trainers/daegc_trainer.py
@@ -0,0 +1,124 @@
+from tqdm import tqdm
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.sparse
+from sklearn.cluster import KMeans
+from sklearn.metrics.cluster import normalized_mutual_info_score
+
+from cogdl.utils import add_remaining_self_loops, spmm
+from .base_trainer import BaseTrainer
+
+
+class DAEGCTrainer(BaseTrainer):
+    def __init__(self, args):
+        self.num_clusters = args.num_clusters
+        self.max_epoch = args.max_epoch
+        self.lr = args.lr
+        self.weight_decay = args.weight_decay
+        self.T = args.T
+        self.gamma = args.gamma
+        self.device = args.device_id[0] if not args.cpu else "cpu"
+
+    @staticmethod
+    def build_trainer_from_args(args):
+        pass
+
+    def fit(self, model, data):
+        # edge_index_2hop = model.get_2hop(data.edge_index)
+        data.adj = torch.sparse_coo_tensor(
+            data.edge_index, torch.ones(data.edge_index.shape[1]), torch.Size([data.x.shape[0], data.x.shape[0]])
+        ).to_dense()
+        data.adj += torch.eye(data.x.shape[0])
+        data.apply(lambda x: x.to(self.device))
+        edge_index_2hop = data.edge_index
+        model = model.to(self.device)
+        self.num_nodes = data.x.shape[0]
+
+        print("Training initial embedding...")
+        epoch_iter = tqdm(range(self.max_epoch))
+        optimizer = torch.optim.Adam(model.parameters(), lr=self.lr, weight_decay=self.weight_decay)
+
+        for epoch in epoch_iter:
+            model.train()
+            optimizer.zero_grad()
+            z = model(data.x, edge_index_2hop)
+            # print(z)
+            loss = model.recon_loss(z, data.adj)
+            loss.backward()
+            optimizer.step()
+            epoch_iter.set_description(f"Epoch: {epoch:03d}, Loss: {loss:.4f}")
+        print("Getting cluster centers...")
+        kmeans = KMeans(n_clusters=self.num_clusters, random_state=0).fit(
+            model(data.x, data.edge_index).detach().cpu().numpy()
+        )
+        model.cluster_center = torch.nn.Parameter(torch.tensor(kmeans.cluster_centers_, device=self.device))
+
+        print("Self-optimizing...")
+        epoch_iter = tqdm(range(self.max_epoch))
+        # optimizer = torch.optim.Adam(model.parameters(), lr=0.001, weight_decay=self.weight_decay)
+        for epoch in epoch_iter:
+            self.cluster_center = model.cluster_center
+            # print("center=", model.cluster_center)
+            model.train()
+            optimizer.zero_grad()
+            z = model(data.x, edge_index_2hop)
+            # print("z=", z)
+            Q = self.getQ(z)
+            if epoch % self.T == 0:
+                P = self.getP(Q).detach()
+            # print(self.cluster_loss(P, Q))
+            loss = model.recon_loss(z, data.adj) + self.gamma * self.cluster_loss(P, Q)
+            loss.backward()
+            optimizer.step()
+            epoch_iter.set_description(f"Epoch: {epoch:03d}, Loss: {loss:.4f}")
+        return model
+
+    def getQ(self, z):
+        Q = None
+        for i in range(z.shape[0]):
+            dis = torch.sum((z[i].repeat(self.num_clusters, 1) - self.cluster_center) ** 2, dim=1)
+            t = 1 / (1 + dis)
+            t = t / torch.sum(t)
+            if Q is None:
+                Q = t.clone().unsqueeze(0)
+            else:
+                Q = torch.cat((Q, t.unsqueeze(0)), 0)
+        # print("Q=", Q)
+        return Q
+
+    def getP(self, Q):
+        P = torch.sum(Q, dim=0).repeat(Q.shape[0], 1)
+        P = Q ** 2 / P
+        P = P / (torch.ones(1, self.num_clusters, device=self.device) * torch.sum(P, dim=1).unsqueeze(-1))
+        # print("P=", P)
+        return P
+
+    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
+        )
diff --git a/cogdl/trainers/ppr_trainer.py b/cogdl/trainers/ppr_trainer.py
index 054a5933..51468242 100644
--- a/cogdl/trainers/ppr_trainer.py
+++ b/cogdl/trainers/ppr_trainer.py
@@ -28,51 +28,43 @@ def __init__(self, args):
         self.dataset_name = args.dataset
         self.loss_func = None
         self.evaluator = None
-        self.post_process = torch.nn.Identity()
+        self.nprop_inference = args.nprop_inference
 
         self.device = "cpu" if not torch.cuda.is_available() or args.cpu else args.device_id[0]
         self.ppr_norm = args.ppr_norm if hasattr(args, "ppr_norm") else "sym"
 
-    def preprocess_data(self, dataset):
+    def ppr_run(self, dataset, mode="train"):
         data = dataset[0]
         num_nodes = data.x.shape[0]
         nodes = torch.arange(num_nodes)
-        self.train_idx = nodes[data.train_mask]
-        self.test_idx = nodes[data.test_mask]
-        self.val_idx = nodes[data.val_mask]
-        if len(data.y.shape) == 1:
-            self.post_process = torch.nn.LogSoftmax()
+        if mode == "train":
+            mask = data.train_mask
+        elif mode == "val":
+            mask = data.val_mask
+        else:
+            mask = data.test_mask
+        index = nodes[mask].numpy()
 
-        if hasattr(data, "edge_index_train"):
+        if mode == "train" and hasattr(data, "edge_index_train"):
             edge_index = data.edge_index_train
         else:
             edge_index = data.edge_index
 
         if not os.path.exists("./saved"):
             os.mkdir("saved")
-        train_path = f"./saved/{self.dataset_name}_{self.topk}_{self.alpha}_{self.normalization}.train.npz"
-        val_path = f"./saved/{self.dataset_name}_{self.topk}_{self.alpha}_{self.normalization}.val.npz"
+        path = f"./saved/{self.dataset_name}_{self.topk}_{self.alpha}_{self.normalization}.{mode}.npz"
 
-        if os.path.exists(train_path):
-            print("Load Train from cached")
-            train_topk_matrix = sp.load_npz(train_path)
+        if os.path.exists(path):
+            print(f"Load {mode} from cached")
+            train_topk_matrix = sp.load_npz(path)
         else:
+            print(f"Fail to load {mode}")
             train_topk_matrix = build_topk_ppr_matrix_from_data(
-                edge_index, self.alpha, self.epsilon, self.train_idx.numpy(), self.topk, self.normalization
-            )
-            sp.save_npz(train_path, train_topk_matrix)
-        train_dataset = PPRGoDataset(data.x, train_topk_matrix, self.train_idx, data.y)
-
-        if os.path.exists(val_path):
-            print("Load Val from cached")
-            val_topk_matrix = sp.load_npz(val_path)
-        else:
-            val_topk_matrix = build_topk_ppr_matrix_from_data(
-                data.edge_index, self.alpha, self.epsilon, self.val_idx.numpy(), self.topk, self.normalization
+                edge_index, self.alpha, self.epsilon, index, self.topk, self.normalization
             )
-            sp.save_npz(val_path, val_topk_matrix)
-        val_dataset = PPRGoDataset(data.x, val_topk_matrix, self.val_idx, data.y)
-        return train_dataset, val_dataset
+            sp.save_npz(path, train_topk_matrix)
+        result = PPRGoDataset(data.x, train_topk_matrix, index, data.y)
+        return result
 
     def get_dataloader(self, dataset):
         data_loader = torch.utils.data.DataLoader(
@@ -87,12 +79,12 @@ def get_dataloader(self, dataset):
         return data_loader
 
     def fit(self, model, dataset):
-        train_dataset, val_dataset = self.preprocess_data(dataset)
         self.loss_func, self.evaluator = dataset.get_evaluator()
         self.model = model.to(self.device)
         self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.lr, weight_decay=self.weight_decay)
-        train_loader = self.get_dataloader(train_dataset)
-        val_loader = self.get_dataloader(val_dataset)
+
+        train_loader = self.get_dataloader(self.ppr_run(dataset, "train"))
+        val_loader = self.get_dataloader(self.ppr_run(dataset, "val"))
 
         best_loss = 1000
         val_loss = 1000
@@ -111,8 +103,17 @@ def fit(self, model, dataset):
             epoch_iter.set_description(
                 f"Epoch: {epoch}, TrainLoss: {train_loss: .4f}, ValLoss: {val_loss: .4f}, ValAcc: {best_acc: .4f}"
             )
+            print()
         self.model = best_model
-        test_acc = self._test_step(dataset[0])
+
+        del train_loader
+        del val_loader
+        if self.nprop_inference <= 0 or self.dataset_name == "ogbn-papers100M":
+            test_loader = self.get_dataloader(self.ppr_run(dataset, "test"))
+            test_acc, test_loss = self._train_step(test_loader, False)
+        else:
+            test_acc = self._test_step(dataset[0])
+
         print(f"TestAcc: {test_acc: .4f}")
         return dict(Acc=test_acc)
 
@@ -128,7 +129,9 @@ def _train_step(self, loader, is_train=True):
             x, targets, ppr_scores, y = [item.to(self.device) for item in batch]
             if is_train:
                 pred = self.model(x, targets, ppr_scores)
-                pred = self.post_process(pred)
+                if len(y.shape) == 1:
+                    y = y.long()
+                    pred = torch.nn.functional.log_softmax(pred, dim=-1)
                 loss = self.loss_func(pred, y)
                 if len(loss.shape) > 1:
                     loss = torch.sum(torch.mean(loss, dim=0))
@@ -140,7 +143,9 @@ def _train_step(self, loader, is_train=True):
             else:
                 with torch.no_grad():
                     pred = self.model(x, targets, ppr_scores)
-                    pred = self.post_process(pred)
+                    if len(y.shape) == 1:
+                        y = y.long()
+                        pred = torch.nn.functional.log_softmax(pred, dim=-1)
                     loss = self.loss_func(pred, y)
                     if len(loss.shape) > 1:
                         loss = torch.sum(torch.mean(loss, dim=0))
diff --git a/cogdl/trainers/self_task.py b/cogdl/trainers/self_task.py
new file mode 100644
index 00000000..e1c156eb
--- /dev/null
+++ b/cogdl/trainers/self_task.py
@@ -0,0 +1,300 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import scipy.sparse as sp
+import numpy as np
+import networkx as nx
+import random
+
+
+class SSLTask:
+    def __init__(self, edge_index, features, device):
+        self.edge_index = edge_index
+        self.num_nodes = torch.max(edge_index).cpu() + 1
+        self.num_edges = edge_index.shape[1]
+        self.features = features
+        self.device = device
+        self.cached_edges = None
+
+    def transform_data(self):
+        raise NotImplementedError
+
+    def make_loss(self, embeddings):
+        raise NotImplementedError
+
+
+class EdgeMask(SSLTask):
+    def __init__(self, edge_index, features, hidden_size, device):
+        super().__init__(edge_index, features, device)
+        self.linear = nn.Linear(hidden_size, 2).to(device)
+
+    def transform_data(self, mask_ratio=0.1):
+        if self.cached_edges is None:
+            edges = self.edge_index.t()
+            perm = np.random.permutation(self.num_edges)
+            preserve_nnz = int(self.num_edges * (1 - mask_ratio))
+            masked = perm[preserve_nnz:]
+            preserved = perm[:preserve_nnz]
+            self.masked_edges = edges[masked].t()
+            self.cached_edges = edges[preserved].t()
+            mask_num = len(masked)
+            self.neg_edges = self.neg_sample(mask_num)
+            self.pseudo_labels = torch.cat([torch.ones(mask_num), torch.zeros(mask_num)]).long().to(self.device)
+            self.node_pairs = torch.cat([self.masked_edges, self.neg_edges], 1)
+
+        return self.cached_edges, self.features
+
+    def make_loss(self, embeddings):
+        embeddings = self.linear(torch.abs(embeddings[self.node_pairs[0]] - embeddings[self.node_pairs[1]]))
+        output = F.log_softmax(embeddings, dim=1)
+        return F.nll_loss(output, self.pseudo_labels)
+
+    def neg_sample(self, edge_num):
+        edges = self.edge_index.t().cpu().numpy()
+        exclude = set([(_[0], _[1]) for _ in list(edges)])
+        itr = self.sample(exclude)
+        sampled = [next(itr) for _ in range(edge_num)]
+        return torch.tensor(sampled).t().to(self.device)
+
+    def sample(self, exclude):
+        while True:
+            t = tuple(np.random.randint(0, self.num_nodes, 2))
+            if t[0] != t[1] and t not in exclude:
+                exclude.add(t)
+                exclude.add((t[1], t[0]))
+                yield t
+
+
+class PairwiseDistance(SSLTask):
+    def __init__(self, edge_index, features, hidden_size, num_class, device):
+        super().__init__(edge_index, features, device)
+        self.linear = nn.Linear(hidden_size, num_class).to(device)
+        self.nclass = num_class
+        self.get_distance()
+
+    def get_distance(self):
+        G = nx.Graph()
+        G.add_edges_from(self.edge_index.cpu().t().numpy())
+
+        path_length = dict(nx.all_pairs_shortest_path_length(G, cutoff=self.nclass - 1))
+        distance = -np.ones((self.num_nodes, self.num_nodes)).astype(int)
+        for u, p in path_length.items():
+            for v, d in p.items():
+                distance[u][v] = d - 1
+        distance[distance == -1] = self.nclass - 1
+        self.distance = distance
+
+        self.dis_node_pairs = []
+        for i in range(self.nclass):
+            tmp = np.array(np.where(distance == i)).transpose()
+            self.dis_node_pairs.append(tmp)
+
+    def transform_data(self):
+        return self.edge_index, self.features
+
+    def make_loss(self, embeddings, k=4000):
+        node_pairs, pseudo_labels = self.sample(k)
+        # print(node_pairs, pseudo_labels)
+        embeddings = self.linear(torch.abs(embeddings[node_pairs[0]] - embeddings[node_pairs[1]]))
+        output = F.log_softmax(embeddings, dim=1)
+        return F.nll_loss(output, pseudo_labels)
+
+    def sample(self, k):
+        sampled = torch.tensor([]).long()
+        pseudo_labels = torch.tensor([]).long()
+        for i in range(self.nclass):
+            tmp = self.dis_node_pairs[i]
+            indices = np.random.choice(np.arange(len(tmp)), k, replace=False)
+            sampled = torch.cat([sampled, torch.tensor(tmp[indices]).t()], 1)
+            pseudo_labels = torch.cat([pseudo_labels, torch.ones(k).long() * i])
+        return sampled.to(self.device), pseudo_labels.to(self.device)
+
+
+class Distance2Clusters(SSLTask):
+    def __init__(self, edge_index, features, hidden_size, num_clusters, device):
+        super().__init__(edge_index, features, device)
+        self.num_clusters = num_clusters
+        self.linear = nn.Linear(hidden_size, num_clusters).to(device)
+        self.gen_cluster_info()
+
+    def transform_data(self):
+        return self.edge_index, self.features
+
+    def gen_cluster_info(self, use_metis=False):
+        G = nx.Graph()
+        G.add_edges_from(self.edge_index.cpu().t().numpy())
+        if use_metis:
+            import metis
+
+            _, parts = metis.part_graph(G, self.num_clusters)
+        else:
+            from sklearn.cluster import SpectralClustering
+
+            clustering = SpectralClustering(
+                n_clusters=self.num_clusters, assign_labels="discretize", random_state=0
+            ).fit(self.features.cpu())
+            parts = clustering.labels_
+
+        node_clusters = [[] for i in range(self.num_clusters)]
+        for i, p in enumerate(parts):
+            node_clusters[p].append(i)
+        self.central_nodes = np.array([])
+        self.distance_vec = np.zeros((self.num_nodes, self.num_clusters))
+        for i in range(self.num_clusters):
+            subgraph = G.subgraph(node_clusters[i])
+            center = None
+            # print(subgraph.nodes)
+            for node in subgraph.nodes:
+                if center is None or subgraph.degree[node] > subgraph.degree[center]:
+                    center = node
+            np.append(self.central_nodes, center)
+            distance = dict(nx.shortest_path_length(G, source=center))
+            for node in distance:
+                self.distance_vec[node][i] = distance[node]
+        self.distance_vec = torch.tensor(self.distance_vec).float().to(self.device)
+
+    def make_loss(self, embeddings):
+        output = self.linear(embeddings)
+        return F.mse_loss(output, self.distance_vec, reduction="mean")
+
+
+class PairwiseAttrSim(SSLTask):
+    def __init__(self, edge_index, features, hidden_size, k, device):
+        super().__init__(edge_index, features, device)
+        self.k = k
+        self.linear = nn.Linear(hidden_size, 1).to(self.device)
+        self.get_attr_sim()
+
+    def get_avg_distance(self, idx_sorted, k, sampled):
+        self.G = nx.Graph()
+        self.G.add_edges_from(self.edge_index.cpu().t().numpy())
+        avg_min = 0
+        avg_max = 0
+        avg_sampled = 0
+        for i in range(self.num_nodes):
+            distance = dict(nx.shortest_path_length(self.G, source=i))
+            sum = 0
+            num = 0
+            for node in idx_sorted[i, :k]:
+                if node in distance:
+                    sum += distance[node]
+                    num += 1
+            if num:
+                avg_min += sum / num / self.num_nodes.item()
+            sum = 0
+            num = 0
+            for node in idx_sorted[i, -k - 1 :]:
+                if node in distance:
+                    sum += distance[node]
+                    num += 1
+            if num:
+                avg_max += sum / num / self.num_nodes.item()
+            sum = 0
+            num = 0
+            for node in idx_sorted[i, sampled]:
+                if node in distance:
+                    sum += distance[node]
+                    num += 1
+            if num:
+                avg_sampled += sum / num / self.num_nodes.item()
+        return avg_min, avg_max, avg_sampled
+
+    def get_attr_sim(self):
+        from sklearn.metrics.pairwise import cosine_similarity
+
+        sims = cosine_similarity(self.features.detach().cpu().numpy())
+        idx_sorted = sims.argsort(1)
+        self.node_pairs = None
+        self.pseudo_labels = None
+        sampled = self.sample(self.k, self.num_nodes)
+        for i in range(self.num_nodes):
+            for node in np.hstack((idx_sorted[i, : self.k], idx_sorted[i, -self.k - 1 :], idx_sorted[i, sampled])):
+                pair = torch.tensor([[i, node]])
+                sim = torch.tensor([sims[i][node]])
+                self.node_pairs = pair if self.node_pairs is None else torch.cat([self.node_pairs, pair], 0)
+                self.pseudo_labels = sim if self.pseudo_labels is None else torch.cat([self.pseudo_labels, sim])
+        print(
+            "max k avg distance: {%.4f}, min k avg distance: {%.4f}, sampled k avg distance: {%.4f}"
+            % (self.get_avg_distance(idx_sorted, self.k, sampled))
+        )
+        # print(self.node_pairs, self.pseudo_labels)
+        self.node_pairs = self.node_pairs.long().to(self.device)
+        self.pseudo_labels = self.pseudo_labels.float().to(self.device)
+
+    def sample(self, k, num_nodes):
+        sampled = []
+        for i in range(k):
+            sampled.append(int(random.random() * (self.num_nodes - self.k * 2)) + self.k)
+        return np.array(sampled)
+
+    def transform_data(self):
+        return self.edge_index, self.features
+
+    def make_loss(self, embeddings):
+        """
+        k = 5000
+        sampled = np.random.choice(self.node_pairs.shape[0], k, replace=False)
+        node_pairs = self.node_pairs[sampled].t()
+        """
+        node_pairs = self.node_pairs
+        output = self.linear(torch.abs(embeddings[node_pairs[0]] - embeddings[node_pairs[1]]))
+        return F.mse_loss(output, self.pseudo_labels, reduction="mean")
+
+
+class Distance2ClustersPP(SSLTask):
+    def __init__(self, edge_index, features, labels, hidden_size, num_clusters, k, device):
+        super().__init__(edge_index, features, device)
+        self.labels = labels.cpu().numpy()
+        self.k = k
+        self.num_clusters = num_clusters
+        self.clusters = None
+        self.linear = nn.Linear(hidden_size, 1).to(self.device)
+
+    def build_graph(self):
+        edges = self.edge_index.detach().cpu().numpy()
+        edge_attr = np.ones(edges.shape[1])
+        inter_label = np.where(self.labels[edges[0]] - self.labels[edges[1]] != 0)
+        inter_cluster = np.where(self.clusters[edges[0]] - self.clusters[edges[1]] != 0)
+        edge_attr[inter_label] = 2
+        edge_attr[inter_cluster] = 2
+        self.G = nx.Graph()
+        for i in range(edges.shape[1]):
+            self.G.add_edge(edges[0][i], edges[1][i], weight=edge_attr[i])
+
+    def update_cluster(self):
+        if self.clusters is None:
+            from sklearn.cluster import SpectralClustering
+
+            clustering = SpectralClustering(
+                n_clusters=self.num_clusters, assign_labels="discretize", random_state=0
+            ).fit(self.features.cpu())
+            self.clusters = clustering.labels_
+
+        self.build_graph()
+        self.node_pairs = None
+        self.pseudo_labels = None
+        for i in range(self.num_clusters):
+            cluster_idx = np.where(self.clusters == i)[0]
+            if len(cluster_idx) < self.k:
+                continue
+            sampled = np.random.choice(cluster_idx, self.k, replace=False)
+            for node in sampled:
+                distance = dict(nx.shortest_path_length(self.G, source=node))
+                for j in range(self.num_nodes):
+                    if j == node or j not in distance:
+                        continue
+                    pair = torch.tensor([[j, node]])
+                    sim = torch.tensor([distance[j]])
+                    self.node_pairs = pair if self.node_pairs is None else torch.cat([self.node_pairs, pair], 0)
+                    self.pseudo_labels = sim if self.pseudo_labels is None else torch.cat([self.pseudo_labels, sim])
+        self.node_pairs = self.node_pairs.t().to(self.device)
+        self.pseudo_labels = self.pseudo_labels.float().to(self.device)
+
+    def transform_data(self):
+        return self.edge_index, self.features
+
+    def make_loss(self, embeddings):
+        embeddings = self.linear(embeddings)
+        # print(embeddings.shape, self.pseudo_labels.shape, F.mse_loss(embeddings, self.pseudo_labels, reduction="mean"))
+        return F.mse_loss(embeddings, self.pseudo_labels, reduction="mean")
diff --git a/cogdl/trainers/self_task_trainer.py b/cogdl/trainers/self_task_trainer.py
new file mode 100644
index 00000000..9e462233
--- /dev/null
+++ b/cogdl/trainers/self_task_trainer.py
@@ -0,0 +1,127 @@
+from abc import ABC, abstractmethod
+from typing import Any
+import copy
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from tqdm import tqdm
+
+from cogdl.data import Dataset
+from cogdl.models.supervised_model import (
+    SupervisedHeterogeneousNodeClassificationModel,
+    SupervisedHomogeneousNodeClassificationModel,
+)
+from cogdl.trainers.supervised_trainer import SupervisedHomogeneousNodeClassificationTrainer
+from cogdl.trainers.self_task import EdgeMask, PairwiseDistance, Distance2Clusters, PairwiseAttrSim, Distance2ClustersPP
+
+
+class SelfTaskTrainer(SupervisedHomogeneousNodeClassificationTrainer):
+    def __init__(self, args):
+        self.device = args.device_id[0] if not args.cpu else "cpu"
+        self.patience = args.patience
+        self.max_epoch = args.max_epoch
+        self.lr = args.lr
+        self.weight_decay = args.weight_decay
+        self.auxiliary_task = args.auxiliary_task
+        self.hidden_size = args.hidden_size
+        self.alpha = args.alpha
+        self.label_mask = args.label_mask
+
+    @staticmethod
+    def build_trainer_from_args(args):
+        pass
+
+    def resplit_data(self, data):
+        trained = torch.where(data.train_mask)[0]
+        perm = np.random.permutation(trained.shape[0])
+        preserve_nnz = int(len(perm) * (1 - self.label_mask))
+        preserved = trained[perm[:preserve_nnz]]
+        masked = trained[perm[preserve_nnz:]]
+        data.train_mask = torch.zeros(data.train_mask.shape[0], dtype=bool)
+        data.train_mask[preserved] = True
+        data.test_mask[masked] = True
+
+    def fit(self, model: SupervisedHeterogeneousNodeClassificationModel, dataset: Dataset):
+        # self.resplit_data(dataset.data)
+        self.data = dataset.data
+        self.original_data = dataset.data
+        self.data.apply(lambda x: x.to(self.device))
+        self.original_data.apply(lambda x: x.to(self.device))
+        if self.auxiliary_task == "edgemask":
+            self.agent = EdgeMask(self.data.edge_index, self.data.x, self.hidden_size, self.device)
+        elif self.auxiliary_task == "pairwise-distance":
+            self.agent = PairwiseDistance(self.data.edge_index, self.data.x, self.hidden_size, 3, self.device)
+        elif self.auxiliary_task == "distance2clusters":
+            self.agent = Distance2Clusters(self.data.edge_index, self.data.x, self.hidden_size, 30, self.device)
+        elif self.auxiliary_task == "pairwise-attr-sim":
+            self.agent = PairwiseAttrSim(self.data.edge_index, self.data.x, self.hidden_size, 5, self.device)
+        elif self.auxiliary_task == "distance2clusters++":
+            self.agent = Distance2ClustersPP(
+                self.data.edge_index, self.data.x, self.data.y, self.hidden_size, 5, 1, self.device
+            )
+        else:
+            raise Exception(
+                "auxiliary task must be edgemask, pairwise-distance, distance2clusters, distance2clusters++ or pairwise-attr-sim"
+            )
+        self.model = model
+
+        self.optimizer = torch.optim.Adam(
+            list(model.parameters()) + list(self.agent.linear.parameters()), lr=self.lr, weight_decay=self.weight_decay
+        )
+        self.model.to(self.device)
+        epoch_iter = tqdm(range(self.max_epoch))
+
+        best_score = 0
+        best_loss = np.inf
+        max_score = 0
+        min_loss = np.inf
+
+        for epoch in epoch_iter:
+            if self.auxiliary_task == "distance2clusters++" and epoch % 40 == 0:
+                self.agent.update_cluster()
+            elif self.auxiliary_task == "scalable-distance-pred" and epoch % 10 == 0:
+                self.agent.update()
+            self._train_step()
+            train_acc, _ = self._test_step(split="train")
+            val_acc, val_loss = self._test_step(split="val")
+            test_acc, test_loss = self._test_step(split="test")
+            epoch_iter.set_description(
+                f"Epoch: {epoch:03d}, Train: {train_acc:.4f}, Val: {val_acc:.4f}, Test: {test_acc:.4f}"
+            )
+            if val_loss <= min_loss or val_acc >= max_score:
+                if val_loss <= best_loss:  # and val_acc >= best_score:
+                    best_loss = val_loss
+                    best_score = val_acc
+                    best_model = copy.deepcopy(self.model)
+                min_loss = np.min((min_loss, val_loss))
+                max_score = np.max((max_score, val_acc))
+        print(f"Valid accurracy = {best_score}")
+
+        return best_model
+
+    def _train_step(self):
+        self.data.edge_index, self.data.x = self.agent.transform_data()
+        self.model.train()
+        self.optimizer.zero_grad()
+        embeddings = self.model.get_embeddings(self.data.x, self.data.edge_index)
+        loss = self.model.node_classification_loss(self.data) + self.alpha * self.agent.make_loss(embeddings)
+        loss.backward()
+        self.optimizer.step()
+
+    def _test_step(self, split="val"):
+        self.data = self.original_data
+        self.model.eval()
+        if split == "train":
+            mask = self.data.train_mask
+        elif split == "val":
+            mask = self.data.val_mask
+        else:
+            mask = self.data.test_mask
+
+        logits = self.model.predict(self.data)
+        loss = F.nll_loss(logits[mask], self.data.y[mask]).item()
+
+        pred = logits[mask].max(1)[1]
+        acc = pred.eq(self.data.y[mask]).sum().item() / mask.sum().item()
+        return acc, loss
diff --git a/docs/source/models.agc.rst b/docs/source/models.agc.rst
new file mode 100644
index 00000000..3f8044b4
--- /dev/null
+++ b/docs/source/models.agc.rst
@@ -0,0 +1,30 @@
+models.agc package
+=================
+
+Submodules
+----------
+
+models.agc.agc module
+------------------------
+
+.. automodule:: models.agc.agc
+    :members:
+    :undoc-members:
+    :show-inheritance:
+
+models.agc.daegc module
+--------------------
+
+.. automodule:: models.agc.daegc
+    :members:
+    :undoc-members:
+    :show-inheritance:
+
+
+Module contents
+---------------
+
+.. automodule:: models.agc
+    :members:
+    :undoc-members:
+    :show-inheritance:


Test patch:
diff --git a/cogdl/datasets/test_data.py b/cogdl/datasets/test_data.py
new file mode 100644
index 00000000..a3e6db21
--- /dev/null
+++ b/cogdl/datasets/test_data.py
@@ -0,0 +1,23 @@
+import torch
+
+from cogdl.datasets import register_dataset
+from cogdl.data import Dataset, Data
+
+@register_dataset("test_small")
+class TestSmallDataset(Dataset):
+    r"""small dataset for debug"""
+    def __init__(self, args=None):
+        x = torch.FloatTensor([[-2, -1], [-2, 1], [-1, 0], [0, 0], [0, 1], [1, 0], [2, 1], [3, 0], [2, -1], [4, 0], [4, 1], [5, 0]])
+        edge_index = torch.LongTensor([[0, 0, 1, 1, 2, 2, 2, 3, 3, 3, 4, 5, 5, 5, 5, 6, 6, 7, 7, 7, 7, 8, 8, 9, 9 , 9 , 10, 10, 11, 11],
+                                       [1, 2, 0, 2, 0, 1, 3, 2, 4, 5, 3, 3, 6, 7, 8, 5, 7, 5, 6, 8, 9, 5, 7, 7, 10, 11, 9 , 11, 9 , 10]])
+        y = torch.LongTensor([0, 0, 0, 1, 1, 2, 2, 2, 2, 3, 3, 3])
+        self.data = Data(x, edge_index, None, y, None)
+        self.data.train_mask = torch.tensor([True, False, False, True, False, True, False, False, False, True, False, False])
+        self.data.val_mask = torch.tensor([False, True, False, False, False, False, True, False, False, False, False, True])
+        self.data.test_mask = torch.tensor([False, False, True, False, True, False, False, True, True, False, True, False])
+        self.num_classes = 4
+        self.transform = None
+
+    def get(self, idx):
+        assert idx == 0
+        return self.data
\ No newline at end of file
diff --git a/tests/tasks/test_attributed_graph_clustering.py b/tests/tasks/test_attributed_graph_clustering.py
index fb5bdb5e..b295958f 100644
--- a/tests/tasks/test_attributed_graph_clustering.py
+++ b/tests/tasks/test_attributed_graph_clustering.py
@@ -2,7 +2,6 @@
 import torch
 from cogdl.tasks import build_task
 from cogdl.utils import build_args_from_dict
-import unittest
 
 graph_clustering_task_name = "attributed_graph_clustering"
 
@@ -10,15 +9,35 @@
 def get_default_args():
     cuda_available = torch.cuda.is_available()
     default_dict = {
+        "task": graph_clustering_task_name,
         "device_id": [0],
         "num_clusters": 7,
         "cluster_method": "kmeans",
         "hidden_size": 16,
-        "model_type": "emb",
-        "momentum": 0,
+        "model_type": "spectral",
         "enhance": None,
         "cpu": not cuda_available,
+        "step": 5,
+        "theta": 0.5,
+        "mu": 0.2,
         "checkpoint": False,
+        "walk_length": 80,
+        "walk_num": 40,
+        "window_size": 5,
+        "worker": 10,
+        "iteration": 3,
+        "rank": 256,
+        "negative": 1,
+        "is_large": False,
+        "max_iter": 5,
+        "embedding_size": 16,
+        "weight_decay": 0.01,
+        "num_heads": 1,
+        "dropout": 0,
+        "max_epoch": 3,
+        "lr": 0.001,
+        "T": 5,
+        "gamma": 10,
     }
     return build_args_from_dict(default_dict)
 
@@ -34,40 +53,117 @@ def create_simple_task():
     return AttributedGraphClustering(args)
 
 
-class TestAttributedGraphClustering(unittest.TestCase):
-    simple_task = None  # prone cora task
-    simple_task_train_result = None
+def test_kmeans_cora():
+    args = get_default_args()
+    args.model_type = "content"
+    args.model = "prone"
+    args.dataset = "cora"
+    args.cluster_method = "kmeans"
+    task = build_task(args)
+    ret = task.train()
+    assert ret["Accuracy"] > 0
 
-    @classmethod
-    def setUpClass(cls) -> None:
-        if not cls.simple_task:
-            cls.simple_task = create_simple_task()
-            cls.simple_task_train_result = cls.simple_task.train()
 
-    def test_accuracy_within_range(self):
-        assert 0 <= self.__class__.simple_task_train_result["Accuracy"] <= 1
+def test_kmeans_citeseer():
+    args = get_default_args()
+    args.model_type = "content"
+    args.model = "prone"
+    args.dataset = "citeseer"
+    args.cluster_method = "kmeans"
+    args.num_clusters = 6
+    task = build_task(args)
+    ret = task.train()
+    assert ret["Accuracy"] > 0
 
-    def test_micro_f1_within_range(self):
-        assert 0 <= self.__class__.simple_task_train_result["Micro_F1"] <= 1
 
-    def test_nmi_within_range(self):
-        assert 0 <= self.__class__.simple_task_train_result["NMI"] <= 1
+def test_spectral_cora():
+    args = get_default_args()
+    args.model_type = "content"
+    args.model = "prone"
+    args.dataset = "cora"
+    args.cluster_method = "spectral"
+    task = build_task(args)
+    ret = task.train()
+    assert ret["Accuracy"] > 0
 
-    def test_all_metrics_present(self):
-        assert "Accuracy" in self.__class__.simple_task_train_result
-        assert "NMI" in self.__class__.simple_task_train_result
-        assert "Micro_F1" in self.__class__.simple_task_train_result
 
-    def test_correct_class_resolved_from_build_task(self):
-        args = get_default_args()
-        args.task = graph_clustering_task_name
-        args.dataset = "cora"
-        args.model = "prone"
-        args.step = 5
-        args.theta = 0.5
-        args.mu = 0.2
-        assert isinstance(build_task(args), AttributedGraphClustering)
+def test_spectral_citeseer():
+    args = get_default_args()
+    args.model_type = "content"
+    args.model = "prone"
+    args.dataset = "citeseer"
+    args.cluster_method = "spectral"
+    args.num_clusters = 6
+    task = build_task(args)
+    ret = task.train()
+    assert ret["Accuracy"] > 0
+
+
+def test_prone_cora():
+    args = get_default_args()
+    args.model = "prone"
+    args.model_type = "spectral"
+    args.dataset = "cora"
+    args.cluster_method = "kmeans"
+    task = build_task(args)
+    ret = task.train()
+    assert ret["Accuracy"] > 0
+
+
+def test_deepwalk_cora():
+    args = get_default_args()
+    args.model = "deepwalk"
+    args.model_type = "spectral"
+    args.dataset = "cora"
+    args.cluster_method = "kmeans"
+    task = build_task(args)
+    ret = task.train()
+    assert ret["Accuracy"] > 0
+
+
+def test_netmf_cora():
+    args = get_default_args()
+    args.model = "netmf"
+    args.model_type = "spectral"
+    args.dataset = "cora"
+    args.cluster_method = "kmeans"
+    task = build_task(args)
+    ret = task.train()
+    assert ret["Accuracy"] > 0
+
+
+def test_agc_cora():
+    args = get_default_args()
+    args.model = "agc"
+    args.model_type = "both"
+    args.dataset = "cora"
+    args.cluster_method = "spectral"
+    args.max_iter = 10
+    task = build_task(args)
+    ret = task.train()
+    assert ret["Accuracy"] > 0
+
+
+def test_daegc_cora():
+    args = get_default_args()
+    args.model = "daegc"
+    args.model_type = "both"
+    args.dataset = "cora"
+    args.cluster_method = "kmeans"
+    task = build_task(args)
+    ret = task.train()
+    assert ret["Accuracy"] > 0
 
 
 if __name__ == "__main__":
-    unittest.main()
+    test_kmeans_cora()
+    # test_kmeans_citeseer()
+    test_spectral_cora()
+    # test_spectral_citeseer()
+
+    test_agc_cora()
+    test_daegc_cora()
+
+    test_prone_cora()
+    test_deepwalk_cora()
+    test_netmf_cora()
diff --git a/tests/tasks/test_node_classification.py b/tests/tasks/test_node_classification.py
index 00bc93a6..9c60b7c3 100644
--- a/tests/tasks/test_node_classification.py
+++ b/tests/tasks/test_node_classification.py
@@ -208,6 +208,7 @@ def test_pyg_gcn_cora():
     args.task = "node_classification"
     args.dataset = "cora"
     args.model = "pyg_gcn"
+    args.auxiliary_task = "none"
     args.num_layers = 2
     task = build_task(args)
     ret = task.train()
diff --git a/tests/tasks/test_self_task.py b/tests/tasks/test_self_task.py
new file mode 100644
index 00000000..8fdadd3e
--- /dev/null
+++ b/tests/tasks/test_self_task.py
@@ -0,0 +1,98 @@
+import torch
+import random
+import numpy as np
+
+from cogdl.tasks import build_task
+from cogdl.datasets import build_dataset
+from cogdl.models import build_model
+from cogdl.utils import build_args_from_dict
+
+
+def get_default_args():
+    cuda_available = torch.cuda.is_available()
+    args = {
+        "dataset": "cora",
+        "model": "pyg_gcn",
+        "hidden_size": 64,
+        "dropout": 0.5,
+        "patience": 2,
+        "device_id": [0],
+        "max_epoch": 3,
+        "sampler": "none",
+        "cpu": not cuda_available,
+        "lr": 0.01,
+        "weight_decay": 5e-4,
+        "missing_rate": -1,
+        "task": "node_classification",
+        "checkpoint": False,
+        "label_mask": 0,
+        "num_layers": 2,
+    }
+    args = build_args_from_dict(args)
+    dataset = build_dataset(args)
+    args.num_features = dataset.num_features
+    args.num_classes = dataset.num_classes
+    return args
+
+
+def test_edgemask():
+    args = get_default_args()
+    args.auxiliary_task = "edgemask"
+    args.alpha = 1
+    dataset = build_dataset(args)
+    model = build_model(args)
+    task = build_task(args, dataset=dataset, model=model)
+    ret = task.train()
+    assert 0 <= ret["Acc"] <= 1
+
+
+def test_pairwise_distance():
+    args = get_default_args()
+    args.auxiliary_task = "pairwise-distance"
+    args.alpha = 35
+    dataset = build_dataset(args)
+    model = build_model(args)
+    task = build_task(args, dataset=dataset, model=model)
+    ret = task.train()
+    assert 0 <= ret["Acc"] <= 1
+
+
+def test_distance_to_clusters():
+    args = get_default_args()
+    args.auxiliary_task = "distance2clusters"
+    args.alpha = 3
+    dataset = build_dataset(args)
+    model = build_model(args)
+    task = build_task(args, dataset=dataset, model=model)
+    ret = task.train()
+    assert 0 <= ret["Acc"] <= 1
+
+
+def test_pairwise_attr_sim():
+    args = get_default_args()
+    args.auxiliary_task = "pairwise-attr-sim"
+    args.alpha = 100
+    dataset = build_dataset(args)
+    model = build_model(args)
+    task = build_task(args, dataset=dataset, model=model)
+    ret = task.train()
+    assert 0 <= ret["Acc"] <= 1
+
+
+def test_distance_to_clustersPP():
+    args = get_default_args()
+    args.auxiliary_task = "distance2clusters++"
+    args.alpha = 1
+    dataset = build_dataset(args)
+    model = build_model(args)
+    task = build_task(args, dataset=dataset, model=model)
+    ret = task.train()
+    assert 0 <= ret["Acc"] <= 1
+
+
+if __name__ == "__main__":
+    test_edgemask()
+    test_pairwise_distance()
+    test_distance_to_clusters()
+    test_pairwise_attr_sim()
+    test_distance_to_clustersPP()