Repository: THUDM/CogDL (language: python)

Issue:
[Framework] Support DDP training for cluster_dw and graphsage_dw
## Description
<!-- Brief description. Refer to the related issues if existed.-->

## Checklist
Please feel free to remove inapplicable items for your PR.
- [ ] The PR title starts with [$CATEGORY] (such as [Model], [Doc], [Feature], [Bugfix])
- [ ] Changes are complete (i.e. I finished coding on this PR)
- [ ] All changes have test coverage
- [ ] Code is well-documented
- [ ] To my best knowledge, examples are either not affected by this change,
      or have been fixed to be compatible with this change
- [ ] Related issue is referred in this PR



Discussion / hints:

Fix patch:
diff --git a/cogdl/__init__.py b/cogdl/__init__.py
index 23f75a46..4d399c34 100644
--- a/cogdl/__init__.py
+++ b/cogdl/__init__.py
@@ -1,5 +1,5 @@
 __version__ = "0.4.1"
 
-from .experiments import experiment
+from .runner.experiments import experiment
 from .oag import oagbert
 from .pipelines import pipeline
diff --git a/cogdl/data/__init__.py b/cogdl/data/__init__.py
index 0d14f999..550ca06d 100644
--- a/cogdl/data/__init__.py
+++ b/cogdl/data/__init__.py
@@ -1,6 +1,6 @@
 from .data import Graph, Adjacency
-from .batch import Batch
+from .batch import Batch, batch_graphs
 from .dataset import Dataset, MultiGraphDataset
 from .dataloader import DataLoader
 
-__all__ = ["Graph", "Adjacency", "Batch", "Dataset", "DataLoader", "MultiGraphDataset"]
+__all__ = ["Graph", "Adjacency", "Batch", "Dataset", "DataLoader", "MultiGraphDataset", "batch_graphs"]
diff --git a/cogdl/data/batch.py b/cogdl/data/batch.py
index 3940d870..217930b2 100644
--- a/cogdl/data/batch.py
+++ b/cogdl/data/batch.py
@@ -4,6 +4,10 @@
 from cogdl.data import Graph, Adjacency
 
 
+def batch_graphs(graphs):
+    return Batch.from_data_list(graphs, class_type=Graph)
+
+
 class Batch(Graph):
     r"""A plain old python object modeling a batch of graphs as one big
     (dicconnected) graph. With :class:`cogdl.data.Data` being the
@@ -19,7 +23,7 @@ def __init__(self, batch=None, **kwargs):
         self.__slices__ = None
 
     @staticmethod
-    def from_data_list(data_list):
+    def from_data_list(data_list, class_type=None):
         r"""Constructs a batch object from a python list holding
         :class:`cogdl.data.Data` objects.
         The assignment vector :obj:`batch` is created on the fly.
@@ -31,8 +35,11 @@ def from_data_list(data_list):
         keys = list(set.union(*keys))
         assert "batch" not in keys
 
-        batch = Batch()
-        batch.__data_class__ = data_list[0].__class__
+        if class_type is not None:
+            batch = class_type()
+        else:
+            batch = Batch()
+            batch.__data_class__ = data_list[0].__class__
         batch.__slices__ = {key: [0] for key in keys}
 
         for key in keys:
diff --git a/cogdl/data/data.py b/cogdl/data/data.py
index c7f716fb..9422a393 100644
--- a/cogdl/data/data.py
+++ b/cogdl/data/data.py
@@ -2,6 +2,7 @@
 import copy
 from contextlib import contextmanager
 import scipy.sparse as sp
+import networkx as nx
 
 import torch
 import numpy as np
@@ -45,7 +46,8 @@ def keys(self):
 
     def __len__(self):
         r"""Returns the number of all present attributes."""
-        return len(self.keys)
+        # return len(self.keys)
+        return 1
 
     def __contains__(self, key):
         r"""Returns :obj:`True`, if the attribute :obj:`key` is present in the
@@ -275,10 +277,12 @@ def set_symmetric(self, val):
         assert val in [True, False]
         self.__symmetric__ = val
 
-    @property
-    def degrees(self):
+    def degrees(self, node_idx=None):
         if self.row_ptr is not None:
-            return (self.row_ptr[1:] - self.row_ptr[:-1]).float()
+            degs = (self.row_ptr[1:] - self.row_ptr[:-1]).float()
+            if node_idx is not None:
+                return degs[node_idx]
+            return degs
         else:
             return get_degrees(self.row, self.col, num_nodes=self.num_nodes)
 
@@ -322,6 +326,10 @@ def num_nodes(self):
             self.__num_nodes__ = max(self.row.max().item(), self.col.max().item()) + 1
             return self.__num_nodes__
 
+    @property
+    def row_ptr_v(self):
+        return self.row_ptr
+
     @property
     def device(self):
         return self.row.device if self.row is not None else self.row_ptr.device
@@ -397,14 +405,25 @@ def to_scipy_csr(self):
             mx = sp.csr_matrix((data, col_ind, row_ptr), shape=(num_nodes, num_nodes))
         return mx
 
-    def random_walk(self, start, length=1, restart_p=0.0):
+    def to_networkx(self, weighted=True):
+        gnx = nx.Graph()
+        row, col = self.edge_index
+        row = row.tolist()
+        col = col.tolist()
+
+        if weighted:
+            weight = self.get_weight().tolist()
+            gnx.add_weighted_edges_from([(row[i], col[i], weight[i]) for i in range(len(row))])
+        else:
+            edges = torch.stack((row, col)).cpu().numpy().transpose()
+            gnx.add_edges_from(edges)
+        return gnx
+
+    def random_walk(self, seeds, length=1, restart_p=0.0):
         if not hasattr(self, "__walker__"):
             scipy_adj = self.to_scipy_csr()
             self.__walker__ = RandomWalker(scipy_adj)
-        return self.__walker__.walk(start, length, restart_p=restart_p)
-
-    def random_walk_with_restart(self, start, length, restart_p):
-        return self.random_walk(start, length, restart_p)
+        return self.__walker__.walk(seeds, length, restart_p=restart_p)
 
     @staticmethod
     def from_dict(dictionary):
@@ -478,6 +497,7 @@ def __init__(self, x=None, y=None, **kwargs):
         self._adj = self._adj_full
         self.__is_train__ = False
         self.__temp_adj_stack__ = list()
+        self.__temp_storage__ = dict()
 
     def train(self):
         self.__is_train__ = True
@@ -604,14 +624,12 @@ def edge_types(self, edge_types):
 
     @property
     def row_indptr(self):
-        if self._adj.row_ptr is None:
-            self._adj.convert_csr()
-        return self._adj.row_ptr
+        return self._adj.row_indptr
 
     @property
     def col_indices(self):
         if self._adj.row_ptr is None:
-            self._adj.convert_csr()
+            self._adj._to_csr()
         return self._adj.col
 
     @row_indptr.setter
@@ -636,7 +654,7 @@ def keys(self):
         keys = [key for key in keys if key[:2] != "__" and key[-2:] != "__"]
         return keys
 
-    def degrees(self):
+    def degrees(self, node_idx=None):
         return self._adj.degrees
 
     def __keys__(self):
@@ -707,6 +725,20 @@ def from_pyg_data(data):
     def clone(self):
         return Graph.from_dict({k: v.clone() for k, v in self})
 
+    def store(self, key):
+        if hasattr(self, key) and not callable(getattr(self, key)):
+            self.__temp_storage__[key] = copy.deepcopy(getattr(self, key))
+        if hasattr(self._adj, key) and not callable(getattr(self._adj, key)):
+            self.__temp_storage__[key] = copy.deepcopy(getattr(self._adj, key))
+
+    def restore(self, key):
+        if key in self.__temp_storage__:
+            if hasattr(self, key) and not callable(getattr(self, key)):
+                setattr(self, key, self.__temp_storage__[key])
+            elif hasattr(self._adj, key) and not callable(getattr(self._adj, key)):
+                self(self._adj, key, self.__temp_storage__[key])
+            self.__temp_storage__.pop(key)
+
     def __delitem__(self, key):
         if hasattr(self, key):
             self[key] = None
@@ -720,7 +752,10 @@ def sample_adj(self, batch, size=-1, replace=True):
         if sample_adj_c is not None:
             if not torch.is_tensor(batch):
                 batch = torch.tensor(batch, dtype=torch.long)
-            (row_ptr, col_indices, nodes, edges) = sample_adj_c(self._adj.row_ptr, self._adj.col, batch, size, replace)
+            # (row_ptr, col_indices, nodes, edges) = sample_adj_c(self._adj.row_ptr, self._adj.col, batch, size, replace)
+            (row_ptr, col_indices, nodes, edges) = sample_adj_c(
+                self._adj.row_indptr, self.col_indices, batch, size, replace
+            )
         else:
             if not (batch[1:] > batch[:-1]).all():
                 batch = batch.sort()[0]
@@ -775,9 +810,15 @@ def _sample_adj(self, batch_size, indices, indptr, size):
         return row_ptr, edge_cols
 
     def csr_subgraph(self, node_idx):
-        if self._adj.row_ptr is None:
+        if self._adj.row_ptr_v is None:
             self._adj._to_csr()
-        indptr, indices, nodes, edges = subgraph_c(self._adj.row_ptr, self._adj.col, node_idx.cpu())
+        if torch.is_tensor(node_idx):
+            node_idx = node_idx.cpu()
+        else:
+            node_idx = torch.as_tensor(node_idx)
+
+        node_idx = torch.unique(node_idx)
+        indptr, indices, nodes, edges = subgraph_c(self._adj.row_ptr, self._adj.col, node_idx)
         nodes_idx = node_idx.to(self._adj.device)
 
         data = Graph(row_ptr=indptr, col=indices)
@@ -838,9 +879,18 @@ def edge_subgraph(self, edge_idx, require_idx=True):
         else:
             return g
 
+    def random_walk(self, seeds, max_nodes_per_seed, restart_p=0.0):
+        return self._adj.random_walk(seeds, max_nodes_per_seed, restart_p)
+
+    def random_walk_with_restart(self, seeds, max_nodes_per_seed, restart_p=0.0):
+        return self._adj.random_walk(seeds, max_nodes_per_seed, restart_p)
+
     def to_scipy_csr(self):
         return self._adj.to_scipy_csr()
 
+    def to_networkx(self):
+        return self._adj.to_networkx()
+
     @staticmethod
     def from_dict(dictionary):
         r"""Creates a data object from a python dictionary."""
@@ -849,6 +899,9 @@ def from_dict(dictionary):
             data[key] = item
         return data
 
+    def nodes(self):
+        return torch.arange(self.num_nodes)
+
     # @property
     # def requires_grad(self):
     #     return False
diff --git a/cogdl/data/dataloader.py b/cogdl/data/dataloader.py
index 0c913d4f..5b25c5b0 100644
--- a/cogdl/data/dataloader.py
+++ b/cogdl/data/dataloader.py
@@ -1,10 +1,18 @@
+from abc import ABCMeta
 import torch.utils.data
 from torch.utils.data.dataloader import default_collate
 
 from cogdl.data import Batch, Graph
 
 
-class DataLoader(torch.utils.data.DataLoader):
+class RecordParameters(ABCMeta):
+    def __call__(cls, *args, **kwargs):
+        obj = type.__call__(cls, *args, **kwargs)
+        obj.record_parameters([args, kwargs])
+        return obj
+
+
+class DataLoader(torch.utils.data.DataLoader, metaclass=RecordParameters):
     r"""Data loader which merges data objects from a
     :class:`cogdl.data.dataset` to a mini-batch.
 
@@ -17,12 +25,13 @@ class DataLoader(torch.utils.data.DataLoader):
     """
 
     def __init__(self, dataset, batch_size=1, shuffle=True, **kwargs):
+        if "collate_fn" not in kwargs or kwargs["collate_fn"] is None:
+            kwargs["collate_fn"] = self.collate_fn
+
         super(DataLoader, self).__init__(
-            # dataset, batch_size, shuffle, collate_fn=lambda data_list: Batch.from_data_list(data_list), **kwargs
             dataset,
             batch_size,
             shuffle,
-            collate_fn=self.collate_fn,
             **kwargs,
         )
 
@@ -37,3 +46,9 @@ def collate_fn(batch):
             return torch.tensor(batch, dtype=torch.float)
 
         raise TypeError("DataLoader found invalid type: {}".format(type(item)))
+
+    def get_parameters(self):
+        return self.default_kwargs
+
+    def record_parameters(self, params):
+        self.default_kwargs = params
diff --git a/cogdl/data/dataset.py b/cogdl/data/dataset.py
index 513f3ad6..785e114d 100644
--- a/cogdl/data/dataset.py
+++ b/cogdl/data/dataset.py
@@ -225,7 +225,7 @@ def get(self, idx):
         if isinstance(idx, slice):
             start = idx.start
             end = idx.stop
-            step = idx.step
+            step = idx.step if idx.step else 1
             idx = list(range(start, end, step))
 
         if len(idx) > 1:
diff --git a/cogdl/data/sampler.py b/cogdl/data/sampler.py
index c3c4936e..6d05165b 100644
--- a/cogdl/data/sampler.py
+++ b/cogdl/data/sampler.py
@@ -7,7 +7,7 @@
 import torch.utils.data
 
 from cogdl.utils import remove_self_loops, row_normalization
-from cogdl.data import Graph
+from cogdl.data import Graph, DataLoader
 
 
 def normalize(adj):
@@ -59,7 +59,7 @@ class SAINTBaseSampler(object):
             ==> Need to override the `preproc()` in sub-class
      2. Post-processing: upon getting the sampled subgraphs, we need to prepare the
             appropriate information (e.g., subgraph adj with renamed indices) to
-            enable the PyTorch trainer.
+            enable the PyTorch controller.
     """
 
     def __init__(self, data, args_params):
@@ -250,7 +250,7 @@ def __getitem__(self, idx):
         return data
 
 
-class SAINTDataLoader(torch.utils.data.DataLoader):
+class SAINTDataLoader(DataLoader):
     def __init__(self, dataset, **kwargs):
         self.dataset = dataset
         kwargs["batch_size"] = 1
@@ -385,7 +385,7 @@ def __call__(self, *args, **kwargs):
         return self.sampler
 
 
-class NeighborSampler(torch.utils.data.DataLoader):
+class NeighborSampler(DataLoader):
     def __init__(self, dataset, sizes: List[int], mask=None, **kwargs):
         if "batch_size" in kwargs:
             batch_size = kwargs["batch_size"]
@@ -524,7 +524,7 @@ def preprocess(self, n_cluster):
         return division
 
 
-class ClusteredLoader(torch.utils.data.DataLoader):
+class ClusteredLoader(DataLoader):
     def __init__(self, dataset, n_cluster: int, method="metis", **kwargs):
         if "batch_size" in kwargs:
             batch_size = kwargs["batch_size"]
diff --git a/cogdl/datasets/customized_data.py b/cogdl/datasets/customized_data.py
index d678be52..ad8edf6b 100644
--- a/cogdl/datasets/customized_data.py
+++ b/cogdl/datasets/customized_data.py
@@ -42,7 +42,7 @@ class NodeDataset(Dataset):
     metric: Accuracy, multi-label f1 or multi-class f1. Default: `accuracy`
     """
 
-    def __init__(self, path="cus_data.pt", scale_feat=True, metric="accuracy"):
+    def __init__(self, path="cus_data.pt", scale_feat=True, metric="auto"):
         self.path = path
         super(NodeDataset, self).__init__(root=path)
         try:
@@ -54,10 +54,11 @@ def __init__(self, path="cus_data.pt", scale_feat=True, metric="accuracy"):
             exit(1)
         self.metric = metric
         if hasattr(self.data, "y") and self.data.y is not None:
-            if len(self.data.y.shape) > 1:
-                self.metric = "multilabel_f1"
-            else:
-                self.metric = "accuracy"
+            if metric == "auto":
+                if len(self.data.y.shape) > 1:
+                    self.metric = "multilabel_f1"
+                else:
+                    self.metric = "accuracy"
 
     def download(self):
         pass
@@ -81,8 +82,7 @@ def _download(self):
     def _process(self):
         if not os.path.exists(self.path):
             data = self.process()
-            if not os.path.exists(self.path):
-                torch.save(data, self.path)
+            torch.save(data, self.path)
 
     def __repr__(self):
         return "{}()".format(self.name)
@@ -112,8 +112,7 @@ def process(self):
     def _process(self):
         if not os.path.exists(self.path):
             data = self.process()
-            if not os.path.exists(self.path):
-                torch.save(data, self.path)
+            torch.save(data, self.path)
 
     def get_evaluator(self):
         return _get_evaluator(self.metric)
diff --git a/cogdl/datasets/gatne.py b/cogdl/datasets/gatne.py
index f87b43a3..d3348053 100644
--- a/cogdl/datasets/gatne.py
+++ b/cogdl/datasets/gatne.py
@@ -82,9 +82,6 @@ def process(self):
     def __repr__(self):
         return "{}()".format(self.name)
 
-    def __len__(self):
-        return self.data.y.shape[0]
-
 
 @register_dataset("amazon")
 class AmazonDataset(GatneDataset):
diff --git a/cogdl/datasets/gtn_data.py b/cogdl/datasets/gtn_data.py
index 52246815..af35b4eb 100644
--- a/cogdl/datasets/gtn_data.py
+++ b/cogdl/datasets/gtn_data.py
@@ -103,6 +103,7 @@ def get(self, idx):
 
     def apply_to_device(self, device):
         self.data.x = self.data.x.to(device)
+        self.data.y = self.data.y.to(device)
 
         self.data.train_node = self.data.train_node.to(device)
         self.data.valid_node = self.data.valid_node.to(device)
diff --git a/cogdl/datasets/han_data.py b/cogdl/datasets/han_data.py
index 694fe4c6..930a3ff8 100644
--- a/cogdl/datasets/han_data.py
+++ b/cogdl/datasets/han_data.py
@@ -103,6 +103,12 @@ def read_gtn_data(self, folder):
         data.test_node = torch.from_numpy(test_idx[0]).type(torch.LongTensor)
         data.test_target = torch.from_numpy(y_test).type(torch.LongTensor)
 
+        y = np.zeros((num_nodes), dtype=int)
+        x_index = torch.cat((data.train_node, data.valid_node, data.test_node))
+        y_index = torch.cat((data.train_target, data.valid_target, data.test_target))
+        y[x_index.numpy()] = y_index.numpy()
+        data.y = torch.from_numpy(y)
+
         self.data = data
 
     def get(self, idx):
@@ -111,6 +117,7 @@ def get(self, idx):
 
     def apply_to_device(self, device):
         self.data.x = self.data.x.to(device)
+        self.data.y = self.data.y.to(device)
 
         self.data.train_node = self.data.train_node.to(device)
         self.data.valid_node = self.data.valid_node.to(device)
diff --git a/cogdl/datasets/ogb.py b/cogdl/datasets/ogb.py
index a4138e1c..b2e1466e 100644
--- a/cogdl/datasets/ogb.py
+++ b/cogdl/datasets/ogb.py
@@ -8,8 +8,7 @@
 
 from . import register_dataset
 from cogdl.data import Dataset, Graph, DataLoader
-from cogdl.utils import cross_entropy_loss, accuracy, remove_self_loops, coalesce, bce_with_logits_loss
-from torch_geometric.utils import to_undirected
+from cogdl.utils import CrossEntropyLoss, Accuracy, remove_self_loops, coalesce, BCEWithLogitsLoss, to_undirected
 
 
 class OGBNDataset(Dataset):
@@ -26,10 +25,10 @@ def get(self, idx):
         return self.data
 
     def get_loss_fn(self):
-        return cross_entropy_loss
+        return CrossEntropyLoss()
 
     def get_evaluator(self):
-        return accuracy
+        return Accuracy()
 
     def _download(self):
         pass
@@ -111,7 +110,7 @@ def edge_attr_size(self):
         ]
 
     def get_loss_fn(self):
-        return bce_with_logits_loss
+        return BCEWithLogitsLoss()
 
     def get_evaluator(self):
         evaluator = NodeEvaluator(name="ogbn-proteins")
@@ -253,8 +252,7 @@ def process(self):
             edge_tps = np.full(src.shape, edge_type_dict[k])
 
             if edge_type == "cites":
-                _edges = torch.as_tensor([src, tgt])
-                _src, _tgt = to_undirected(_edges).numpy()
+                _src, _tgt = to_undirected([src, tgt]).numpy()
                 edge_tps = np.full(_src.shape, edge_type_dict[k])
                 edge_idx = np.vstack([_src, _tgt])
             else:
@@ -343,7 +341,7 @@ def __init__(self, root, name):
         self.name = name
         self.dataset = GraphPropPredDataset(self.name, root)
 
-        self.graphs = []
+        self.data = []
         self.all_nodes = 0
         self.all_edges = 0
         for i in range(len(self.dataset.graphs)):
@@ -355,7 +353,7 @@ def __init__(self, root, name):
                 y=torch.tensor(label),
             )
             data.num_nodes = graph["num_nodes"]
-            self.graphs.append(data)
+            self.data.append(data)
 
             self.all_nodes += graph["num_nodes"]
             self.all_edges += graph["edge_index"].shape[1]
@@ -372,11 +370,11 @@ def get_loader(self, args):
     def get_subset(self, subset):
         datalist = []
         for idx in subset:
-            datalist.append(self.graphs[idx])
+            datalist.append(self.data[idx])
         return datalist
 
     def get(self, idx):
-        return self.graphs[idx]
+        return self.data[idx]
 
     def _download(self):
         pass
diff --git a/cogdl/datasets/planetoid_data.py b/cogdl/datasets/planetoid_data.py
index bde333c7..cc628016 100644
--- a/cogdl/datasets/planetoid_data.py
+++ b/cogdl/datasets/planetoid_data.py
@@ -6,7 +6,7 @@
 import torch
 
 from cogdl.data import Dataset, Graph
-from cogdl.utils import remove_self_loops, download_url, untar, coalesce
+from cogdl.utils import remove_self_loops, download_url, untar, coalesce, Accuracy, CrossEntropyLoss
 from . import register_dataset
 
 
@@ -147,6 +147,11 @@ def num_classes(self):
         assert hasattr(self.data, "y")
         return int(torch.max(self.data.y)) + 1
 
+    @property
+    def num_nodes(self):
+        assert hasattr(self.data, "y")
+        return self.data.y.shape[0]
+
     def download(self):
         fname = "{}.zip".format(self.name.lower())
         download_url("{}{}.zip&dl=1".format(self.url, self.name.lower()), self.raw_dir, fname)
@@ -165,6 +170,12 @@ def __repr__(self):
     def __len__(self):
         return 1
 
+    def get_evaluator(self):
+        return Accuracy()
+
+    def get_loss_fn(self):
+        return CrossEntropyLoss()
+
 
 def normalize_feature(data):
     x_sum = torch.sum(data.x, dim=1)
diff --git a/cogdl/datasets/saint_data.py b/cogdl/datasets/saint_data.py
index da6e6dbe..1e2a4cab 100644
--- a/cogdl/datasets/saint_data.py
+++ b/cogdl/datasets/saint_data.py
@@ -8,7 +8,7 @@
 from sklearn.preprocessing import StandardScaler
 
 from cogdl.data import Graph, Dataset
-from cogdl.utils import download_url, accuracy, multilabel_f1, bce_with_logits_loss, cross_entropy_loss
+from cogdl.utils import download_url, Accuracy, MultiLabelMicroF1, BCEWithLogitsLoss, CrossEntropyLoss
 from . import register_dataset
 from .planetoid_data import index_to_mask
 
@@ -101,10 +101,10 @@ def get(self, idx):
         return self.data
 
     def get_evaluator(self):
-        return multilabel_f1
+        return Accuracy()
 
     def get_loss_fn(self):
-        return bce_with_logits_loss
+        return BCEWithLogitsLoss()
 
     def __repr__(self):
         return "{}()".format(self.name)
@@ -143,10 +143,10 @@ def __init__(self, data_path="data"):
         self.data = scale_feats(self.data)
 
     def get_evaluator(self):
-        return multilabel_f1
+        return MultiLabelMicroF1()
 
     def get_loss_fn(self):
-        return bce_with_logits_loss
+        return BCEWithLogitsLoss()
 
 
 @register_dataset("amazon-s")
@@ -161,10 +161,10 @@ def __init__(self, data_path="data"):
         self.data = scale_feats(self.data)
 
     def get_evaluator(self):
-        return multilabel_f1
+        return MultiLabelMicroF1()
 
     def get_loss_fn(self):
-        return bce_with_logits_loss
+        return BCEWithLogitsLoss()
 
 
 @register_dataset("flickr")
@@ -179,10 +179,10 @@ def __init__(self, data_path="data"):
         self.data = scale_feats(self.data)
 
     def get_evaluator(self):
-        return accuracy
+        return Accuracy()
 
     def get_loss_fn(self):
-        return cross_entropy_loss
+        return CrossEntropyLoss()
 
 
 @register_dataset("reddit")
@@ -197,10 +197,10 @@ def __init__(self, data_path="data"):
         self.data = scale_feats(self.data)
 
     def get_evaluator(self):
-        return accuracy
+        return Accuracy()
 
     def get_loss_fn(self):
-        return cross_entropy_loss
+        return CrossEntropyLoss()
 
 
 @register_dataset("ppi")
@@ -215,10 +215,10 @@ def __init__(self, data_path="data"):
         self.data = scale_feats(self.data)
 
     def get_evaluator(self):
-        return multilabel_f1
+        return MultiLabelMicroF1()
 
     def get_loss_fn(self):
-        return bce_with_logits_loss
+        return BCEWithLogitsLoss()
 
 
 @register_dataset("ppi-large")
@@ -235,7 +235,7 @@ def __init__(self, data_path="data"):
         self.data = scale_feats(self.data)
 
     def get_evaluator(self):
-        return multilabel_f1
+        return MultiLabelMicroF1()
 
     def get_loss_fn(self):
-        return bce_with_logits_loss
+        return BCEWithLogitsLoss()
diff --git a/cogdl/layers/__init__.py b/cogdl/layers/__init__.py
index d55e30c3..f3c03e45 100644
--- a/cogdl/layers/__init__.py
+++ b/cogdl/layers/__init__.py
@@ -16,6 +16,7 @@
 from .sgc_layer import SGCLayer
 from .mixhop_layer import MixHopLayer
 from .reversible_layer import RevGNNLayer
+from .set2set import Set2Set
 
 
 __all__ = [
@@ -40,4 +41,5 @@
     "MixHopLayer",
     "MLP",
     "RevGNNLayer",
+    "Set2Set",
 ]
diff --git a/cogdl/layers/gat_layer.py b/cogdl/layers/gat_layer.py
index 542374b4..f1438138 100644
--- a/cogdl/layers/gat_layer.py
+++ b/cogdl/layers/gat_layer.py
@@ -13,26 +13,26 @@ class GATLayer(nn.Module):
     """
 
     def __init__(
-        self, in_features, out_features, nhead=1, alpha=0.2, attn_drop=0.5, activation=None, residual=False, norm=None
+        self, in_feats, out_feats, nhead=1, alpha=0.2, attn_drop=0.5, activation=None, residual=False, norm=None
     ):
         super(GATLayer, self).__init__()
-        self.in_features = in_features
-        self.out_features = out_features
+        self.in_features = in_feats
+        self.out_features = out_feats
         self.alpha = alpha
         self.nhead = nhead
 
-        self.W = nn.Parameter(torch.FloatTensor(in_features, out_features * nhead))
+        self.W = nn.Parameter(torch.FloatTensor(in_feats, out_feats * nhead))
 
-        self.a_l = nn.Parameter(torch.zeros(size=(1, nhead, out_features)))
-        self.a_r = nn.Parameter(torch.zeros(size=(1, nhead, out_features)))
+        self.a_l = nn.Parameter(torch.zeros(size=(1, nhead, out_feats)))
+        self.a_r = nn.Parameter(torch.zeros(size=(1, nhead, out_feats)))
 
         self.dropout = nn.Dropout(attn_drop)
         self.leakyrelu = nn.LeakyReLU(self.alpha)
         self.act = None if activation is None else get_activation(activation)
-        self.norm = None if norm is None else get_norm_layer(norm, out_features * nhead)
+        self.norm = None if norm is None else get_norm_layer(norm, out_feats * nhead)
 
         if residual:
-            self.residual = nn.Linear(in_features, out_features * nhead)
+            self.residual = nn.Linear(in_feats, out_feats * nhead)
         else:
             self.register_buffer("residual", None)
         self.reset_parameters()
diff --git a/cogdl/layers/gcn_layer.py b/cogdl/layers/gcn_layer.py
index 03902e69..41783fd3 100644
--- a/cogdl/layers/gcn_layer.py
+++ b/cogdl/layers/gcn_layer.py
@@ -53,7 +53,7 @@ def forward(self, graph, x):
         if self.norm is not None:
             out = self.norm(out)
         if self.act is not None:
-            out = self.act(out, inplace=True)
+            out = self.act(out)
 
         if self.residual is not None:
             out = out + self.residual(x)
diff --git a/cogdl/layers/mlp_layer.py b/cogdl/layers/mlp_layer.py
index 1fc5a571..05b2c289 100644
--- a/cogdl/layers/mlp_layer.py
+++ b/cogdl/layers/mlp_layer.py
@@ -63,12 +63,12 @@ def forward(self, x):
         for i, fc in enumerate(self.mlp[:-1]):
             x = fc(x)
             if self.act_first:
-                x = self.activation(x, inplace=True)
+                x = self.activation(x)
             if self.norm:
                 x = self.norm_list[i](x)
 
             if not self.act_first:
-                x = self.activation(x, inplace=True)
+                x = self.activation(x)
             x = F.dropout(x, p=self.dropout, training=self.training)
         x = self.mlp[-1](x)
         return x
diff --git a/cogdl/layers/sage_layer.py b/cogdl/layers/sage_layer.py
index 7ff9b366..579b3e7c 100644
--- a/cogdl/layers/sage_layer.py
+++ b/cogdl/layers/sage_layer.py
@@ -19,7 +19,9 @@ def __call__(self, graph, x):
 
 
 class SAGELayer(nn.Module):
-    def __init__(self, in_feats, out_feats, normalize=False, aggr="mean", dropout=0.0, norm=None, activation=None):
+    def __init__(
+        self, in_feats, out_feats, normalize=False, aggr="mean", dropout=0.0, norm=None, activation=None, residual=False
+    ):
         super(SAGELayer, self).__init__()
         self.in_feats = in_feats
         self.out_feats = out_feats
@@ -51,6 +53,11 @@ def __init__(self, in_feats, out_feats, normalize=False, aggr="mean", dropout=0.
         else:
             self.norm = None
 
+        if residual:
+            self.residual = nn.Linear(in_features=in_feats, out_features=out_feats)
+        else:
+            self.residual = None
+
     def forward(self, graph, x):
         out = self.aggr(graph, x)
         out = torch.cat([x, out], dim=-1)
@@ -63,6 +70,10 @@ def forward(self, graph, x):
         if self.act is not None:
             out = self.act(out, inplace=True)
 
+        if self.residual:
+            out = out + self.residual(x)
+
         if self.dropout is not None:
             out = self.dropout(out)
+
         return out
diff --git a/cogdl/layers/set2set.py b/cogdl/layers/set2set.py
new file mode 100644
index 00000000..b2506b68
--- /dev/null
+++ b/cogdl/layers/set2set.py
@@ -0,0 +1,63 @@
+import torch
+from cogdl.utils import mul_edge_softmax, batch_sum_pooling
+
+
+class Set2Set(torch.nn.Module):
+    r"""The global pooling operator based on iterative content-based attention
+    from the `"Order Matters: Sequence to sequence for sets"
+    <https://arxiv.org/abs/1511.06391>`_ paper
+
+    .. math::
+        \mathbf{q}_t &= \mathrm{LSTM}(\mathbf{q}^{*}_{t-1})
+
+        \alpha_{i,t} &= \mathrm{softmax}(\mathbf{x}_i \cdot \mathbf{q}_t)
+
+        \mathbf{r}_t &= \sum_{i=1}^N \alpha_{i,t} \mathbf{x}_i
+
+        \mathbf{q}^{*}_t &= \mathbf{q}_t \, \Vert \, \mathbf{r}_t,
+
+    where :math:`\mathbf{q}^{*}_T` defines the output of the layer with twice
+    the dimensionality as the input.
+
+    Args:
+        in_channels (int): Size of each input sample.
+        processing_steps (int): Number of iterations :math:`T`.
+        num_layers (int, optional): Number of recurrent layers, *.e.g*, setting
+            :obj:`num_layers=2` would mean stacking two LSTMs together to form
+            a stacked LSTM, with the second LSTM taking in outputs of the first
+            LSTM and computing the final results. (default: :obj:`1`)
+    """
+
+    def __init__(self, in_feats, processing_steps, num_layers=1):
+        super(Set2Set, self).__init__()
+
+        self.in_channels = in_feats
+        self.out_channels = 2 * in_feats
+        self.processing_steps = processing_steps
+        self.num_layers = num_layers
+
+        self.lstm = torch.nn.LSTM(self.out_channels, self.in_channels, num_layers)
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        self.lstm.reset_parameters()
+
+    def forward(self, x, batch):
+        batch_size = batch.max().item() + 1
+
+        h = (
+            x.new_zeros((self.num_layers, batch_size, self.in_channels)),
+            x.new_zeros((self.num_layers, batch_size, self.in_channels)),
+        )
+        q_star = x.new_zeros(batch_size, self.out_channels)
+
+        for i in range(self.processing_steps):
+            q, h = self.lstm(q_star.unsqueeze(0), h)
+            q = q.view(batch_size, self.in_channels)
+            e = (x * q[batch]).sum(dim=-1, keepdim=True)
+            a = mul_edge_softmax(e, batch)
+            r = batch_sum_pooling(a * x, batch)
+            q_star = torch.cat([q, r], dim=-1)
+
+        return q_star
diff --git a/cogdl/loggers/__init__.py b/cogdl/loggers/__init__.py
new file mode 100644
index 00000000..6382e0b1
--- /dev/null
+++ b/cogdl/loggers/__init__.py
@@ -0,0 +1,14 @@
+from .base_logger import Logger
+
+
+def build_logger(logger, log_path="./runs", project="cogdl-exp"):
+    if logger == "wandb":
+        from .wandb_logger import WandbLogger
+
+        return WandbLogger(log_path, project)
+    elif logger == "tensorboard":
+        from .tensorboard_logger import TBLogger
+
+        return TBLogger(log_path)
+    else:
+        return Logger(log_path)
diff --git a/cogdl/loggers/base_logger.py b/cogdl/loggers/base_logger.py
new file mode 100644
index 00000000..28fd9c30
--- /dev/null
+++ b/cogdl/loggers/base_logger.py
@@ -0,0 +1,12 @@
+class Logger:
+    def __init__(self, log_path):
+        self.log_path = log_path
+
+    def start(self):
+        pass
+
+    def note(self, metrics, step=None):
+        pass
+
+    def finish(self):
+        pass
diff --git a/cogdl/loggers/tensorboard_logger.py b/cogdl/loggers/tensorboard_logger.py
new file mode 100644
index 00000000..74771c7a
--- /dev/null
+++ b/cogdl/loggers/tensorboard_logger.py
@@ -0,0 +1,24 @@
+from tensorboardX import SummaryWriter
+
+from . import Logger
+
+
+class TBLogger(Logger):
+    def __init__(self, log_path):
+        super(TBLogger, self).__init__(log_path)
+        self.last_step = 0
+
+    def start(self):
+        self.writer = SummaryWriter(logdir=self.log_path)
+
+    def note(self, metrics, step=None):
+        if not hasattr(self, "writer"):
+            self.start()
+        if step is None:
+            step = self.last_step
+        for key, value in metrics.items():
+            self.writer.add_scalar(key, value, step)
+        self.last_step = step
+
+    def finish(self):
+        self.writer.close()
diff --git a/cogdl/loggers/wandb_logger.py b/cogdl/loggers/wandb_logger.py
new file mode 100644
index 00000000..5a639326
--- /dev/null
+++ b/cogdl/loggers/wandb_logger.py
@@ -0,0 +1,28 @@
+import warnings
+from . import Logger
+
+try:
+    import wandb
+except Exception:
+    warnings.warn("Please install wandb first")
+
+
+class WandbLogger(Logger):
+    def __init__(self, log_path, project=None):
+        super(WandbLogger, self).__init__(log_path)
+        self.last_step = 0
+        self.project = project
+
+    def start(self):
+        self.run = wandb.init(reinit=True, dir=self.log_path, project=self.project)
+
+    def note(self, metrics, step=None):
+        if not hasattr(self, "run"):
+            self.start()
+        if step is None:
+            step = self.last_step
+        self.run.log(metrics, step=step)
+        self.last_step = step
+
+    def finish(self):
+        self.run.finish()
diff --git a/cogdl/match.yml b/cogdl/match.yml
index 22aa245f..c729d76f 100644
--- a/cogdl/match.yml
+++ b/cogdl/match.yml
@@ -186,13 +186,11 @@ multiplex_node_classification:
 heterogeneous_node_classification:
   - model:
     - gtn
+    - han
     dataset:
     - gtn-dblp
     - gtn-acm
     - gtn-imdb
-  - model:
-    - han
-    dataset:
     - han-dblp
     - han-acm
     - han-imdb
diff --git a/cogdl/models/__init__.py b/cogdl/models/__init__.py
index c62ad8f5..d79d85db 100644
--- a/cogdl/models/__init__.py
+++ b/cogdl/models/__init__.py
@@ -121,7 +121,7 @@ def build_model(args):
     "sortpool": "cogdl.models.nn.sortpool",
     "srgcn": "cogdl.models.nn.pyg_srgcn",
     "asgcn": "cogdl.models.nn.asgcn",
-    "gcc": "cogdl.models.nn.dgl_gcc",
+    "gcc": "cogdl.models.nn.gcc_model",
     "unsup_graphsage": "cogdl.models.nn.unsup_graphsage",
     "sagpool": "cogdl.models.nn.pyg_sagpool",
     "graphsaint": "cogdl.models.nn.graphsaint",
diff --git a/cogdl/models/emb/deepwalk.py b/cogdl/models/emb/deepwalk.py
index 8a3ac0ed..65783842 100644
--- a/cogdl/models/emb/deepwalk.py
+++ b/cogdl/models/emb/deepwalk.py
@@ -26,16 +26,17 @@ class DeepWalk(BaseModel):
     def add_args(parser: argparse.ArgumentParser):
         """Add model-specific arguments to the parser."""
         # fmt: off
-        parser.add_argument('--walk-length', type=int, default=80,
-                            help='Length of walk per source. Default is 80.')
-        parser.add_argument('--walk-num', type=int, default=40,
-                            help='Number of walks per source. Default is 40.')
-        parser.add_argument('--window-size', type=int, default=5,
-                            help='Window size of skip-gram model. Default is 5.')
-        parser.add_argument('--worker', type=int, default=10,
-                            help='Number of parallel workers. Default is 10.')
-        parser.add_argument('--iteration', type=int, default=10,
-                            help='Number of iterations. Default is 10.')
+        parser.add_argument("--walk-length", type=int, default=80,
+                            help="Length of walk per source. Default is 80.")
+        parser.add_argument("--walk-num", type=int, default=40,
+                            help="Number of walks per source. Default is 40.")
+        parser.add_argument("--window-size", type=int, default=5,
+                            help="Window size of skip-gram model. Default is 5.")
+        parser.add_argument("--worker", type=int, default=10,
+                            help="Number of parallel workers. Default is 10.")
+        parser.add_argument("--iteration", type=int, default=10,
+                            help="Number of iterations. Default is 10.")
+        parser.add_argument("--hidden-size", type=int, default=128)
         # fmt: on
 
     @classmethod
@@ -58,8 +59,12 @@ def __init__(self, dimension, walk_length, walk_num, window_size, worker, iterat
         self.worker = worker
         self.iteration = iteration
 
-    def train(self, G: nx.Graph, embedding_model_creator=Word2Vec):
-        self.G = G
+    def train(self, graph, embedding_model_creator=Word2Vec, return_dict=False):
+        return self.forward(graph, embedding_model_creator, return_dict)
+
+    def forward(self, graph, embedding_model_creator=Word2Vec, return_dict=False):
+        nx_g = graph.to_networkx()
+        self.G = nx_g
         walks = self._simulate_walks(self.walk_length, self.walk_num)
         walks = [[str(node) for node in walk] for walk in walks]
         print("training word2vec...")
@@ -72,9 +77,18 @@ def train(self, G: nx.Graph, embedding_model_creator=Word2Vec):
             workers=self.worker,
             iter=self.iteration,
         )
-        id2node = dict([(vid, node) for vid, node in enumerate(G.nodes())])
+        id2node = dict([(vid, node) for vid, node in enumerate(nx_g.nodes())])
         embeddings = np.asarray([model.wv[str(id2node[i])] for i in range(len(id2node))])
-        return embeddings
+
+        if return_dict:
+            features_matrix = dict()
+            for vid, node in enumerate(nx_g.nodes()):
+                features_matrix[node] = embeddings[vid]
+        else:
+            features_matrix = np.zeros((graph.num_nodes, embeddings.shape[1]))
+            nx_nodes = nx_g.nodes()
+            features_matrix[nx_nodes] = embeddings[np.arange(graph.num_nodes)]
+        return features_matrix
 
     def _walk(self, start_node, walk_length):
         # Simulate a random walk starting from start node.
diff --git a/cogdl/models/emb/dngr.py b/cogdl/models/emb/dngr.py
index b5135493..860e4edc 100644
--- a/cogdl/models/emb/dngr.py
+++ b/cogdl/models/emb/dngr.py
@@ -123,7 +123,11 @@ def get_emb(self, matrix):
         emb_matrix = preprocessing.normalize(emb_matrix, "l2")
         return emb_matrix
 
-    def train(self, G):
+    def train(self, graph, return_dict=False):
+        return self.forward(graph, return_dict=False)
+
+    def forward(self, graph, return_dict=False):
+        G = graph.to_networkx()
         self.num_node = G.number_of_nodes()
         A = nx.adjacency_matrix(G).todense()
         PPMI = self.get_ppmi_matrix(A)
@@ -147,5 +151,15 @@ def train(self, G):
             Loss.backward()
             epoch_iter.set_description(f"Epoch: {epoch:03d},  Loss: {Loss:.8f}")
             opt.step()
-        embedding, _ = model.forward(input_mat)
-        return embedding.detach().cpu().numpy()
+        embeddings, _ = model.forward(input_mat)
+        embeddings = embeddings.detach().cpu().numpy()
+
+        if return_dict:
+            features_matrix = dict()
+            for vid, node in enumerate(G.nodes()):
+                features_matrix[node] = embeddings[vid]
+        else:
+            features_matrix = np.zeros((graph.num_nodes, embeddings.shape[1]))
+            nx_nodes = G.nodes()
+            features_matrix[nx_nodes] = embeddings[np.arange(graph.num_nodes)]
+        return features_matrix
diff --git a/cogdl/models/emb/gatne.py b/cogdl/models/emb/gatne.py
index 3a749bee..58db542f 100644
--- a/cogdl/models/emb/gatne.py
+++ b/cogdl/models/emb/gatne.py
@@ -108,8 +108,12 @@ def __init__(
         self.schema = schema
 
         self.multiplicity = True
-
+        self.device = "cpu" if not torch.cuda.is_available() else "cuda"
+    
     def train(self, network_data):
+        return self.forward(network_data)
+
+    def forward(self, network_data):
         all_walks = generate_walks(network_data, self.walk_num, self.walk_length, schema=self.schema)
         vocab, index2word = generate_vocab(all_walks)
         train_pairs = generate_pairs(all_walks, vocab)
diff --git a/cogdl/models/emb/grarep.py b/cogdl/models/emb/grarep.py
index 4272e129..8da219bf 100644
--- a/cogdl/models/emb/grarep.py
+++ b/cogdl/models/emb/grarep.py
@@ -19,8 +19,9 @@ class GraRep(BaseModel):
     def add_args(parser):
         """Add model-specific arguments to the parser."""
         # fmt: off
-        parser.add_argument('--step', type=int, default=5,
-                            help='Number of matrix step in GraRep. Default is 5.')
+        parser.add_argument("--step", type=int, default=5,
+                            help="Number of matrix step in GraRep. Default is 5.")
+        parser.add_argument("--hidden-size", type=int, default=128)
         # fmt: on
 
     @classmethod
@@ -32,9 +33,12 @@ def __init__(self, dimension, step):
         self.dimension = dimension
         self.step = step
 
-    def train(self, G):
-        self.G = G
-        self.num_node = G.number_of_nodes()
+    def train(self, graph, return_dict=False):
+        return self.forward(graph, return_dict)
+
+    def forward(self, graph, return_dict=False):
+        self.G = graph.to_networkx()
+        self.num_node = self.G.number_of_nodes()
         A = np.asarray(nx.adjacency_matrix(self.G).todense(), dtype=float)
         A = preprocessing.normalize(A, "l1")
 
@@ -62,8 +66,16 @@ def train(self, G):
                 W = self._get_embedding(A_list[k], self.dimension / self.step)
                 final_emb = np.hstack((final_emb, W))
 
-        self.embeddings = final_emb
-        return self.embeddings
+        embeddings = final_emb
+        if return_dict:
+            features_matrix = dict()
+            for vid, node in enumerate(self.G.nodes()):
+                features_matrix[node] = embeddings[vid]
+        else:
+            features_matrix = np.zeros((graph.num_nodes, embeddings.shape[1]))
+            nx_nodes = self.G.nodes()
+            features_matrix[nx_nodes] = embeddings[np.arange(graph.num_nodes)]
+        return features_matrix
 
     def _get_embedding(self, matrix, dimension):
         # get embedding from svd and process normalization for ut
diff --git a/cogdl/models/emb/hin2vec.py b/cogdl/models/emb/hin2vec.py
index 1c14d3c2..9e76a514 100644
--- a/cogdl/models/emb/hin2vec.py
+++ b/cogdl/models/emb/hin2vec.py
@@ -173,9 +173,17 @@ def __init__(self, hidden_dim, walk_length, walk_num, batch_size, hop, negative,
         self.epochs = epochs
         self.lr = lr
 
-    def train(self, G, node_type):
+        self.device = "cpu" if not torch.cuda.is_available() or cpu else "cuda"
+
+    def forward(self, data):
+        return self.train(data)
+
+    def train(self, data):
+        G = nx.DiGraph()
+        row, col = data.edge_index
+        G.add_edges_from(list(zip(row.numpy(), col.numpy())))
         self.num_node = G.number_of_nodes()
-        rw = RWgraph(G, node_type)
+        rw = RWgraph(G, data.pos.tolist())
         walks = rw._simulate_walks(self.walk_length, self.walk_num)
         pairs, relation = rw.data_preparation(walks, self.hop, self.negative)
 
diff --git a/cogdl/models/emb/hope.py b/cogdl/models/emb/hope.py
index d3500edf..09cc54bf 100644
--- a/cogdl/models/emb/hope.py
+++ b/cogdl/models/emb/hope.py
@@ -19,8 +19,9 @@ class HOPE(BaseModel):
     def add_args(parser):
         """Add model-specific arguments to the parser."""
         # fmt: off
-        parser.add_argument('--beta', type=float, default=0.01,
-                            help='Parameter of katz for HOPE. Default is 0.01')
+        parser.add_argument("--beta", type=float, default=0.01,
+                            help="Parameter of katz for HOPE. Default is 0.01")
+        parser.add_argument("--hidden-size", type=int, default=128)
         # fmt: on
 
     @classmethod
@@ -32,16 +33,29 @@ def __init__(self, dimension, beta):
         self.dimension = dimension
         self.beta = beta
 
-    def train(self, G):
+    def train(self, graph, return_dict=False):
+        return self.forward(graph, return_dict)
+
+    def forward(self, graph, return_dict=False):
         r"""The author claim that Katz has superior performance in related tasks
         S_katz = (M_g)^-1 * M_l = (I - beta*A)^-1 * beta*A = (I - beta*A)^-1 * (I - (I -beta*A))
         = (I - beta*A)^-1 - I
         """
-        adj = nx.adjacency_matrix(G).todense()
+        nx_g = graph.to_networkx()
+        adj = nx.adjacency_matrix(nx_g).todense()
         n = adj.shape[0]
         katz_matrix = np.asarray((np.eye(n) - self.beta * np.mat(adj)).I - np.eye(n))
-        self.embeddings = self._get_embedding(katz_matrix, self.dimension)
-        return self.embeddings
+        embeddings = self._get_embedding(katz_matrix, self.dimension)
+
+        if return_dict:
+            features_matrix = dict()
+            for vid, node in enumerate(nx_g.nodes()):
+                features_matrix[node] = embeddings[vid]
+        else:
+            features_matrix = np.zeros((graph.num_nodes, embeddings.shape[1]))
+            nx_nodes = nx_g.nodes()
+            features_matrix[nx_nodes] = embeddings[np.arange(graph.num_nodes)]
+        return features_matrix
 
     def _get_embedding(self, matrix, dimension):
         # get embedding from svd and process normalization for ut and vt
diff --git a/cogdl/models/emb/line.py b/cogdl/models/emb/line.py
index 3081eb0e..64d3c192 100644
--- a/cogdl/models/emb/line.py
+++ b/cogdl/models/emb/line.py
@@ -29,18 +29,19 @@ class LINE(BaseModel):
     def add_args(parser):
         """Add model-specific arguments to the parser."""
         # fmt: off
-        parser.add_argument('--walk-length', type=int, default=80,
-                            help='Length of walk per source. Default is 80.')
-        parser.add_argument('--walk-num', type=int, default=20,
-                            help='Number of walks per source. Default is 20.')
-        parser.add_argument('--negative', type=int, default=5,
-                            help='Number of negative node in sampling. Default is 5.')
-        parser.add_argument('--batch-size', type=int, default=1000,
-                            help='Batch size in SGD training process. Default is 1000.')
-        parser.add_argument('--alpha', type=float, default=0.025,
-                            help='Initial learning rate of SGD. Default is 0.025.')
-        parser.add_argument('--order', type=int, default=3,
-                            help='Order of proximity in LINE. Default is 3 for 1+2.')
+        parser.add_argument("--walk-length", type=int, default=80,
+                            help="Length of walk per source. Default is 80.")
+        parser.add_argument("--walk-num", type=int, default=20,
+                            help="Number of walks per source. Default is 20.")
+        parser.add_argument("--negative", type=int, default=5,
+                            help="Number of negative node in sampling. Default is 5.")
+        parser.add_argument("--batch-size", type=int, default=1000,
+                            help="Batch size in SGD training process. Default is 1000.")
+        parser.add_argument("--alpha", type=float, default=0.025,
+                            help="Initial learning rate of SGD. Default is 0.025.")
+        parser.add_argument("--order", type=int, default=3,
+                            help="Order of proximity in LINE. Default is 3 for 1+2.")
+        parser.add_argument("--hidden-size", type=int, default=128)
         # fmt: on
 
     @classmethod
@@ -56,6 +57,7 @@ def build_model_from_args(cls, args):
         )
 
     def __init__(self, dimension, walk_length, walk_num, negative, batch_size, alpha, order):
+        super(LINE, self).__init__()
         self.dimension = dimension
         self.walk_length = walk_length
         self.walk_num = walk_num
@@ -64,25 +66,29 @@ def __init__(self, dimension, walk_length, walk_num, negative, batch_size, alpha
         self.init_alpha = alpha
         self.order = order
 
-    def train(self, G):
+    def train(self, graph, return_dict=False):
+        return self.train(graph, return_dict)
+
+    def forward(self, graph, return_dict=False):
         # run LINE algorithm, 1-order, 2-order or 3(1-order + 2-order)
-        self.G = G
+        nx_g = graph.to_networkx()
+        self.G = nx_g
         self.is_directed = nx.is_directed(self.G)
-        self.num_node = G.number_of_nodes()
-        self.num_edge = G.number_of_edges()
+        self.num_node = nx_g.number_of_nodes()
+        self.num_edge = nx_g.number_of_edges()
         self.num_sampling_edge = self.walk_length * self.walk_num * self.num_node
 
-        node2id = dict([(node, vid) for vid, node in enumerate(G.nodes())])
+        node2id = dict([(node, vid) for vid, node in enumerate(nx_g.nodes())])
         self.edges = [[node2id[e[0]], node2id[e[1]]] for e in self.G.edges()]
-        self.edges_prob = np.asarray([G[u][v].get("weight", 1.0) for u, v in G.edges()])
+        self.edges_prob = np.asarray([nx_g[u][v].get("weight", 1.0) for u, v in nx_g.edges()])
         self.edges_prob /= np.sum(self.edges_prob)
         self.edges_table, self.edges_prob = alias_setup(self.edges_prob)
 
         degree_weight = np.asarray([0] * self.num_node)
-        for u, v in G.edges():
-            degree_weight[node2id[u]] += G[u][v].get("weight", 1.0)
+        for u, v in nx_g.edges():
+            degree_weight[node2id[u]] += nx_g[u][v].get("weight", 1.0)
             if not self.is_directed:
-                degree_weight[node2id[v]] += G[u][v].get("weight", 1.0)
+                degree_weight[node2id[v]] += nx_g[u][v].get("weight", 1.0)
         self.node_prob = np.power(degree_weight, 0.75)
         self.node_prob /= np.sum(self.node_prob)
         self.node_table, self.node_prob = alias_setup(self.node_prob)
@@ -104,13 +110,22 @@ def train(self, G):
             embedding2 = preprocessing.normalize(self.emb_vertex, "l2")
 
         if self.order == 1:
-            self.embeddings = embedding1
+            embeddings = embedding1
         elif self.order == 2:
-            self.embeddings = embedding2
+            embeddings = embedding2
         else:
             print("concatenate two embedding...")
-            self.embeddings = np.hstack((embedding1, embedding2))
-        return self.embeddings
+            embeddings = np.hstack((embedding1, embedding2))
+
+        if return_dict:
+            features_matrix = dict()
+            for vid, node in enumerate(nx_g.nodes()):
+                features_matrix[node] = embeddings[vid]
+        else:
+            features_matrix = np.zeros((graph.num_nodes, embeddings.shape[1]))
+            nx_nodes = nx_g.nodes()
+            features_matrix[nx_nodes] = embeddings[np.arange(graph.num_nodes)]
+        return features_matrix
 
     def _update(self, vec_u, vec_v, vec_error, label):
         # update vetex embedding and vec_error
diff --git a/cogdl/models/emb/metapath2vec.py b/cogdl/models/emb/metapath2vec.py
index 7e6045e6..fd7bb40d 100644
--- a/cogdl/models/emb/metapath2vec.py
+++ b/cogdl/models/emb/metapath2vec.py
@@ -64,9 +64,15 @@ def __init__(self, dimension, walk_length, walk_num, window_size, worker, iterat
         self.schema = schema
         self.node_type = None
 
-    def train(self, G, node_type):
+    def forward(self, data):
+        return self.train(data)
+
+    def train(self, data):
+        G = nx.DiGraph()
+        row, col = data.edge_index
+        G.add_edges_from(list(zip(row.numpy(), col.numpy())))
         self.G = G
-        self.node_type = [str(a) for a in node_type]
+        self.node_type = [str(a) for a in data.pos.tolist()]
         walks = self._simulate_walks(self.walk_length, self.walk_num, self.schema)
         walks = [[str(node) for node in walk] for walk in walks]
         model = Word2Vec(
diff --git a/cogdl/models/emb/netmf.py b/cogdl/models/emb/netmf.py
index f6a4505f..f43c58d1 100644
--- a/cogdl/models/emb/netmf.py
+++ b/cogdl/models/emb/netmf.py
@@ -25,7 +25,8 @@ def add_args(parser):
         parser.add_argument("--window-size", type=int, default=5)
         parser.add_argument("--rank", type=int, default=256)
         parser.add_argument("--negative", type=int, default=1)
-        parser.add_argument('--is-large', action='store_true')
+        parser.add_argument("--is-large", action="store_true")
+        parser.add_argument("--hidden-size", type=int, default=128)
         # fmt: on
 
     @classmethod
@@ -33,14 +34,19 @@ def build_model_from_args(cls, args):
         return cls(args.hidden_size, args.window_size, args.rank, args.negative, args.is_large)
 
     def __init__(self, dimension, window_size, rank, negative, is_large=False):
+        super(NetMF, self).__init__()
         self.dimension = dimension
         self.window_size = window_size
         self.rank = rank
         self.negative = negative
         self.is_large = is_large
 
-    def train(self, G):
-        A = sp.csr_matrix(nx.adjacency_matrix(G))
+    def train(self, graph, return_dict=False):
+        return self.forward(graph, return_dict)
+
+    def forward(self, graph, return_dict=False):
+        nx_g = graph.to_networkx()
+        A = sp.csr_matrix(nx.adjacency_matrix(nx_g))
         if not self.is_large:
             print("Running NetMF for a small window size...")
             deepwalk_matrix = self._compute_deepwalk_matrix(A, window=self.window_size, b=self.negative)
@@ -54,8 +60,17 @@ def train(self, G):
             )
         # factorize deepwalk matrix with SVD
         u, s, _ = sp.linalg.svds(deepwalk_matrix, self.dimension)
-        self.embeddings = sp.diags(np.sqrt(s)).dot(u.T).T
-        return self.embeddings
+        embeddings = sp.diags(np.sqrt(s)).dot(u.T).T
+
+        if return_dict:
+            features_matrix = dict()
+            for vid, node in enumerate(nx_g.nodes()):
+                features_matrix[node] = embeddings[vid]
+        else:
+            features_matrix = np.zeros((graph.num_nodes, embeddings.shape[1]))
+            nx_nodes = nx_g.nodes()
+            features_matrix[nx_nodes] = embeddings[np.arange(graph.num_nodes)]
+        return features_matrix
 
     def _compute_deepwalk_matrix(self, A, window, b):
         # directly compute deepwalk matrix
diff --git a/cogdl/models/emb/netsmf.py b/cogdl/models/emb/netsmf.py
index eaf5c75f..d1ebeae0 100644
--- a/cogdl/models/emb/netsmf.py
+++ b/cogdl/models/emb/netsmf.py
@@ -28,14 +28,15 @@ class NetSMF(BaseModel):
     def add_args(parser):
         """Add model-specific arguments to the parser."""
         # fmt: off
-        parser.add_argument('--window-size', type=int, default=10,
-                            help='Window size of approximate matrix. Default is 10.')
-        parser.add_argument('--negative', type=int, default=1,
-                            help='Number of negative node in sampling. Default is 1.')
-        parser.add_argument('--num-round', type=int, default=100,
-                            help='Number of round in NetSMF. Default is 100.')
-        parser.add_argument('--worker', type=int, default=10,
-                            help='Number of parallel workers. Default is 10.')
+        parser.add_argument("--window-size", type=int, default=10,
+                            help="Window size of approximate matrix. Default is 10.")
+        parser.add_argument("--negative", type=int, default=1,
+                            help="Number of negative node in sampling. Default is 1.")
+        parser.add_argument("--num-round", type=int, default=100,
+                            help="Number of round in NetSMF. Default is 100.")
+        parser.add_argument("--worker", type=int, default=10,
+                            help="Number of parallel workers. Default is 10.")
+        parser.add_argument("--hidden-size", type=int, default=128)
         # fmt: on
 
     @classmethod
@@ -56,12 +57,15 @@ def __init__(self, dimension, window_size, negative, num_round, worker):
         self.worker = worker
         self.num_round = num_round
 
-    def train(self, G):
-        self.G = G
-        node2id = dict([(node, vid) for vid, node in enumerate(G.nodes())])
+    def train(self, graph, return_dict=False):
+        return self.forward(graph, return_dict)
+
+    def forward(self, graph, return_dict=False):
+        self.G = graph.to_networkx()
+        node2id = dict([(node, vid) for vid, node in enumerate(self.G.nodes())])
         self.is_directed = nx.is_directed(self.G)
         self.num_node = self.G.number_of_nodes()
-        self.num_edge = G.number_of_edges()
+        self.num_edge = self.G.number_of_edges()
         self.edges = [[node2id[e[0]], node2id[e[1]]] for e in self.G.edges()]
 
         id2node = dict(zip(node2id.values(), node2id.keys()))
@@ -72,13 +76,13 @@ def train(self, G):
         self.alias_nodes = {}
         self.node_weight = {}
         for i in range(self.num_node):
-            unnormalized_probs = [G[id2node[i]][nbr].get("weight", 1.0) for nbr in G.neighbors(id2node[i])]
+            unnormalized_probs = [self.G[id2node[i]][nbr].get("weight", 1.0) for nbr in self.G.neighbors(id2node[i])]
             norm_const = sum(unnormalized_probs)
             normalized_probs = [float(u_prob) / norm_const for u_prob in unnormalized_probs]
             self.alias_nodes[i] = alias_setup(normalized_probs)
             self.node_weight[i] = dict(
                 zip(
-                    [node2id[nbr] for nbr in G.neighbors(id2node[i])],
+                    [node2id[nbr] for nbr in self.G.neighbors(id2node[i])],
                     unnormalized_probs,
                 )
             )
@@ -118,8 +122,17 @@ def train(self, G):
         print("number of nzz", M.nnz)
         print("construct matrix sparsifier time", time.time() - t2)
 
-        embedding = self._get_embedding_rand(M)
-        return embedding
+        embeddings = self._get_embedding_rand(M)
+
+        if return_dict:
+            features_matrix = dict()
+            for vid, node in enumerate(self.G.nodes()):
+                features_matrix[node] = embeddings[vid]
+        else:
+            features_matrix = np.zeros((graph.num_nodes, embeddings.shape[1]))
+            nx_nodes = self.G.nodes()
+            features_matrix[nx_nodes] = embeddings[np.arange(graph.num_nodes)]
+        return features_matrix
 
     def _get_embedding_rand(self, matrix):
         # Sparse randomized tSVD for fast embedding
diff --git a/cogdl/models/emb/node2vec.py b/cogdl/models/emb/node2vec.py
index 2c2e8261..c1ebbaec 100644
--- a/cogdl/models/emb/node2vec.py
+++ b/cogdl/models/emb/node2vec.py
@@ -43,6 +43,7 @@ def add_args(parser):
                             help='Parameter in node2vec. Default is 1.0.')
         parser.add_argument('--q', type=float, default=1.0,
                             help='Parameter in node2vec. Default is 1.0.')
+        parser.add_argument("--hidden-size", type=int, default=128)
         # fmt: on
 
     @classmethod
@@ -69,7 +70,11 @@ def __init__(self, dimension, walk_length, walk_num, window_size, worker, iterat
         self.p = p
         self.q = q
 
-    def train(self, G):
+    def train(self, graph, return_dict=False):
+        return self.forward(graph, return_dict)
+
+    def forward(self, graph, return_dict=False):
+        G = graph.to_networkx()
         self.G = G
         is_directed = nx.is_directed(self.G)
         for i, j in G.edges():
@@ -89,8 +94,17 @@ def train(self, G):
             iter=self.iteration,
         )
         id2node = dict([(vid, node) for vid, node in enumerate(G.nodes())])
-        self.embeddings = np.asarray([model.wv[str(id2node[i])] for i in range(len(id2node))])
-        return self.embeddings
+        embeddings = np.asarray([model.wv[str(id2node[i])] for i in range(len(id2node))])
+
+        if return_dict:
+            features_matrix = dict()
+            for vid, node in enumerate(G.nodes()):
+                features_matrix[node] = embeddings[vid]
+        else:
+            features_matrix = np.zeros((graph.num_nodes, embeddings.shape[1]))
+            nx_nodes = G.nodes()
+            features_matrix[nx_nodes] = embeddings[np.arange(graph.num_nodes)]
+        return features_matrix
 
     def _node2vec_walk(self, walk_length, start_node):
         # Simulate a random walk starting from start node.
diff --git a/cogdl/models/emb/prone.py b/cogdl/models/emb/prone.py
index ea867f09..99b6c89d 100644
--- a/cogdl/models/emb/prone.py
+++ b/cogdl/models/emb/prone.py
@@ -1,11 +1,12 @@
-import networkx as nx
 import numpy as np
 import scipy.sparse as sp
 from scipy.special import iv
+import networkx as nx
 from sklearn import preprocessing
 from sklearn.utils.extmath import randomized_svd
 
 from cogdl.utils.prone_utils import get_embedding_dense
+from cogdl.data import Graph
 from .. import BaseModel, register_model
 
 
@@ -29,6 +30,7 @@ def add_args(parser):
                             help="Number of items in the chebyshev expansion")
         parser.add_argument("--mu", type=float, default=0.2)
         parser.add_argument("--theta", type=float, default=0.5)
+        parser.add_argument("--hidden-size", type=int, default=128)
         # fmt: on
 
     @classmethod
@@ -42,18 +44,28 @@ def __init__(self, dimension, step, mu, theta):
         self.mu = mu
         self.theta = theta
 
-    def train(self, G):
-        self.num_node = G.number_of_nodes()
+    def train(self, graph, return_dict=False):
+        return self.forward(graph, return_dict)
 
-        self.matrix0 = sp.csr_matrix(nx.adjacency_matrix(G))
+    def forward(self, graph: Graph, return_dict=False):
+        nx_g = graph.to_networkx()
+        self.matrix0 = sp.csr_matrix(nx.adjacency_matrix(nx_g))
 
         features_matrix = self._pre_factorization(self.matrix0, self.matrix0)
 
         embeddings_matrix = self._chebyshev_gaussian(self.matrix0, features_matrix, self.step, self.mu, self.theta)
 
-        self.embeddings = embeddings_matrix
+        embeddings = embeddings_matrix
 
-        return self.embeddings
+        if return_dict:
+            features_matrix = dict()
+            for vid, node in enumerate(nx_g.nodes()):
+                features_matrix[node] = embeddings[vid]
+        else:
+            features_matrix = np.zeros((graph.num_nodes, embeddings.shape[1]))
+            nx_nodes = nx_g.nodes()
+            features_matrix[nx_nodes] = embeddings[np.arange(graph.num_nodes)]
+        return features_matrix
 
     def _get_embedding_rand(self, matrix):
         # Sparse randomized tSVD for fast embedding
diff --git a/cogdl/models/emb/pte.py b/cogdl/models/emb/pte.py
index 29a74396..431c3263 100644
--- a/cogdl/models/emb/pte.py
+++ b/cogdl/models/emb/pte.py
@@ -59,9 +59,15 @@ def __init__(self, dimension, walk_length, walk_num, negative, batch_size, alpha
         self.batch_size = batch_size
         self.init_alpha = alpha
 
-    def train(self, G, node_type):
+    def forward(self, data):
+        return self.train(data)
+
+    def train(self, data):
+        G = nx.DiGraph()
+        row, col = data.edge_index
+        G.add_edges_from(list(zip(row.numpy(), col.numpy())))
         self.G = G
-        self.node_type = node_type
+        self.node_type = data.pos.tolist()
         self.num_node = G.number_of_nodes()
         self.num_edge = G.number_of_edges()
         self.num_sampling_edge = self.walk_length * self.walk_num * self.num_node
diff --git a/cogdl/models/emb/sdne.py b/cogdl/models/emb/sdne.py
index 6bc6961b..daf3931c 100644
--- a/cogdl/models/emb/sdne.py
+++ b/cogdl/models/emb/sdne.py
@@ -108,7 +108,11 @@ def __init__(self, hidden_size1, hidden_size2, droput, alpha, beta, nu1, nu2, ma
         self.lr = lr
         self.cpu = cpu
 
-    def train(self, G):
+    def train(self, graph, return_dict=False):
+        return self.forward(graph, return_dict)
+
+    def forward(self, graph, return_dict=False):
+        G = graph.to_networkx()
         num_node = G.number_of_nodes()
         model = SDNE_layer(
             num_node, self.hidden_size1, self.hidden_size2, self.droput, self.alpha, self.beta, self.nu1, self.nu2
@@ -131,5 +135,15 @@ def train(self, G):
                 f"Epoch: {epoch:03d}, L_1st: {L_1st:.4f}, L_2nd: {L_2nd:.4f}, L_reg: {L_reg:.4f}"
             )
             opt.step()
-        embedding = model.get_emb(A)
-        return embedding.detach().cpu().numpy()
+        embeddings = model.get_emb(A)
+        embeddings = embeddings.detach().cpu().numpy()
+
+        if return_dict:
+            features_matrix = dict()
+            for vid, node in enumerate(G.nodes()):
+                features_matrix[node] = embeddings[vid]
+        else:
+            features_matrix = np.zeros((graph.num_nodes, embeddings.shape[1]))
+            nx_nodes = G.nodes()
+            features_matrix[nx_nodes] = embeddings[np.arange(graph.num_nodes)]
+        return features_matrix
diff --git a/cogdl/models/emb/spectral.py b/cogdl/models/emb/spectral.py
index da11bae3..4db695b0 100644
--- a/cogdl/models/emb/spectral.py
+++ b/cogdl/models/emb/spectral.py
@@ -17,20 +17,35 @@ class Spectral(BaseModel):
     @staticmethod
     def add_args(parser):
         """Add model-specific arguments to the parser."""
-        pass
+        # fmt: off
+        parser.add_argument("--hidden-size", type=int, default=128)
+        # fmt: on
 
     @classmethod
     def build_model_from_args(cls, args):
         return cls(args.hidden_size)
 
-    def __init__(self, dimension):
+    def __init__(self, hidden_size):
         super(Spectral, self).__init__()
-        self.dimension = dimension
+        self.dimension = hidden_size
 
-    def train(self, G):
-        matrix = nx.normalized_laplacian_matrix(G).todense()
+    def train(self, graph, return_dict=False):
+        return self.forward(graph, return_dict)
+
+    def forward(self, graph, return_dict=False):
+        nx_g = graph.to_networkx()
+        matrix = nx.normalized_laplacian_matrix(nx_g).todense()
         matrix = np.eye(matrix.shape[0]) - np.asarray(matrix)
         ut, s, _ = sp.linalg.svds(matrix, self.dimension)
         emb_matrix = ut * np.sqrt(s)
-        emb_matrix = preprocessing.normalize(emb_matrix, "l2")
-        return emb_matrix
+        embeddings = preprocessing.normalize(emb_matrix, "l2")
+
+        if return_dict:
+            features_matrix = dict()
+            for vid, node in enumerate(nx_g.nodes()):
+                features_matrix[node] = embeddings[vid]
+        else:
+            features_matrix = np.zeros((graph.num_nodes, embeddings.shape[1]))
+            nx_nodes = nx_g.nodes()
+            features_matrix[nx_nodes] = embeddings[np.arange(graph.num_nodes)]
+        return features_matrix
diff --git a/cogdl/models/nn/compgcn.py b/cogdl/models/nn/compgcn.py
index a4acb912..58c77c20 100644
--- a/cogdl/models/nn/compgcn.py
+++ b/cogdl/models/nn/compgcn.py
@@ -219,7 +219,6 @@ def add_args(parser):
         parser.add_argument("--num-bases", type=int, default=10)
         parser.add_argument("--num-layers", type=int, default=1)
         parser.add_argument("--sampling-rate", type=float, default=0.01)
-        parser.add_argument("--score-func", type=str, default="conve")
         parser.add_argument("--lbl_smooth", type=float, default=0.1)
         parser.add_argument("--opn", type=str, default="sub")
         # fmt: on
@@ -232,7 +231,6 @@ def build_model_from_args(cls, args):
             hidden_size=args.hidden_size,
             num_bases=args.num_bases,
             sampling_rate=args.sampling_rate,
-            score_func=args.score_func,
             penalty=args.penalty,
             layers=args.num_layers,
             dropout=args.dropout,
@@ -248,14 +246,13 @@ def __init__(
         num_bases=0,
         layers=1,
         sampling_rate=0.01,
-        score_func="conve",
         penalty=0.001,
         dropout=0.0,
         lbl_smooth=0.1,
         opn="sub",
     ):
         BaseModel.__init__(self)
-        GNNLinkPredict.__init__(self, score_func, hidden_size)
+        GNNLinkPredict.__init__(self)
         activation = F.tanh
         self.model = CompGCN(
             num_entities,
@@ -335,15 +332,4 @@ def predict(self, graph):
         indices = torch.arange(0, self.num_entities).to(device)
         x = self.emb(indices)
         node_embed, rel_embed = self.model(graph, x)
-        edge_index, edge_types = graph.edge_index, graph.edge_attr
-        mrr, hits = cal_mrr(
-            node_embed,
-            rel_embed,
-            edge_index,
-            edge_types,
-            scoring=self.scoring,
-            protocol="raw",
-            batch_size=500,
-            hits=[1, 3, 10],
-        )
-        return mrr, hits
+        return node_embed, rel_embed
diff --git a/cogdl/models/nn/daegc.py b/cogdl/models/nn/daegc.py
index ff0d0cb4..33130f25 100644
--- a/cogdl/models/nn/daegc.py
+++ b/cogdl/models/nn/daegc.py
@@ -29,7 +29,7 @@ def add_args(parser):
         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("--T", type=int, default=5)
         parser.add_argument("--gamma", type=float, default=10)
         # fmt: on
 
@@ -49,12 +49,18 @@ def __init__(self, num_features, hidden_size, embedding_size, num_heads, dropout
         self.num_clusters = num_clusters
         self.att1 = GATLayer(num_features, hidden_size, attn_drop=dropout, alpha=0.2, nhead=num_heads)
         self.att2 = GATLayer(hidden_size * num_heads, embedding_size, attn_drop=dropout, alpha=0.2, nhead=1)
-        self.cluster_center = None
+        self.cluster_center = torch.nn.Parameter(torch.FloatTensor(self.num_clusters))
 
     @staticmethod
     def get_trainer(args=None):
         return DAEGCTrainer
 
+    def set_cluster_center(self, center):
+        self.cluster_center.data = torch.tensor(center, device=self.device)
+
+    def get_cluster_center(self):
+        return self.cluster_center.data.detach()
+
     def forward(self, graph):
         x = graph.x
         x = F.dropout(x, p=self.dropout, training=self.training)
diff --git a/cogdl/models/nn/dgl_gcc.py b/cogdl/models/nn/dgl_gcc.py
index 9c40bcd7..f9dea8fa 100644
--- a/cogdl/models/nn/dgl_gcc.py
+++ b/cogdl/models/nn/dgl_gcc.py
@@ -713,7 +713,7 @@ def _convert_idx(self, idx):
         return graph_idx, node_idx
 
 
-@register_model("gcc")
+@register_model("dgl_gcc")
 class GCC(BaseModel):
     @staticmethod
     def add_args(parser):
diff --git a/cogdl/models/nn/dropedge_gcn.py b/cogdl/models/nn/dropedge_gcn.py
index 6160c1fd..1d57bf6e 100644
--- a/cogdl/models/nn/dropedge_gcn.py
+++ b/cogdl/models/nn/dropedge_gcn.py
@@ -551,7 +551,7 @@ def add_args(parser):
                             help="The input layer of the model.")
         parser.add_argument('--outputlayer', default='gcn',
                             help="The output layer of the model.")
-        parser.add_argument('--hidden_size', type=int, default=64,
+        parser.add_argument('--hidden-size', type=int, default=64,
                             help='Number of hidden units.')
         parser.add_argument('--dropout', type=float, default=0.5,
                             help='Dropout rate (1 - keep probability).')
diff --git a/cogdl/models/nn/gcc_model.py b/cogdl/models/nn/gcc_model.py
new file mode 100644
index 00000000..deb3d289
--- /dev/null
+++ b/cogdl/models/nn/gcc_model.py
@@ -0,0 +1,322 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from cogdl.layers import SELayer
+from .. import BaseModel, register_model
+
+from cogdl.layers import MLP, GATLayer, GINLayer
+from cogdl.utils import batch_sum_pooling, batch_mean_pooling, batch_max_pooling
+from cogdl.layers import Set2Set
+
+
+class ApplyNodeFunc(nn.Module):
+    """Update the node feature hv with MLP, BN and ReLU."""
+
+    def __init__(self, mlp, use_selayer):
+        super(ApplyNodeFunc, self).__init__()
+        self.mlp = mlp
+        self.bn = (
+            SELayer(self.mlp.output_dim, int(np.sqrt(self.mlp.output_dim)))
+            if use_selayer
+            else nn.BatchNorm1d(self.mlp.output_dim)
+        )
+
+    def forward(self, h):
+        h = self.mlp(h)
+        h = self.bn(h)
+        h = F.relu(h)
+        return h
+
+
+class GATModel(nn.Module):
+    def __init__(self, in_feats, hidden_size, num_layers, nhead, dropout=0.0, attn_drop=0.0, alpha=0.2, residual=False):
+        super(GATModel, self).__init__()
+        assert hidden_size % nhead == 0
+        self.layers = nn.ModuleList(
+            [
+                GATLayer(
+                    in_feats=in_feats if i > 0 else hidden_size // nhead,
+                    out_feats=hidden_size // nhead,
+                    nhead=nhead,
+                    attn_drop=0.0,
+                    alpha=0.2,
+                    residual=False,
+                    activation=F.leaky_relu if i + 1 < num_layers else None,
+                )
+                for i in range(num_layers)
+            ]
+        )
+
+    def forward(self, graph, x):
+        for i, layer in enumerate(self.layers):
+            x = layer(graph, x)
+        return x
+
+
+class GINModel(nn.Module):
+    def __init__(
+        self,
+        num_layers,
+        in_feats,
+        hidden_dim,
+        out_feats,
+        num_mlp_layers,
+        eps=0,
+        pooling="sum",
+        train_eps=False,
+        dropout=0.5,
+        final_dropout=0.2,
+    ):
+        super(GINModel, self).__init__()
+        self.gin_layers = nn.ModuleList()
+        self.batch_norm = nn.ModuleList()
+        self.num_layers = num_layers
+        for i in range(num_layers - 1):
+            if i == 0:
+                mlp = MLP(in_feats, hidden_dim, hidden_dim, num_mlp_layers, norm="batchnorm")
+            else:
+                mlp = MLP(hidden_dim, hidden_dim, hidden_dim, num_mlp_layers, norm="batchnorm")
+            self.gin_layers.append(GINLayer(mlp, eps, train_eps))
+            self.batch_norm.append(nn.BatchNorm1d(hidden_dim))
+
+        self.linear_prediction = nn.ModuleList()
+        for i in range(self.num_layers):
+            if i == 0:
+                self.linear_prediction.append(nn.Linear(in_feats, out_feats))
+            else:
+                self.linear_prediction.append(nn.Linear(hidden_dim, out_feats))
+        self.dropout = nn.Dropout(dropout)
+
+        if pooling == "sum":
+            self.pool = batch_sum_pooling
+        elif pooling == "mean":
+            self.pool = batch_mean_pooling
+        elif pooling == "max":
+            self.pool = batch_max_pooling
+        else:
+            raise NotImplementedError
+        self.final_drop = nn.Dropout(final_dropout)
+
+    def forward(self, batch, n_feat):
+        h = n_feat
+        # device = h.device
+        # batchsize = int(torch.max(batch.batch)) + 1
+
+        layer_rep = [h]
+        for i in range(self.num_layers - 1):
+            h = self.gin_layers[i](batch, h)
+            h = self.batch_norm[i](h)
+            h = F.relu(h)
+            layer_rep.append(h)
+
+        score_over_layer = 0
+
+        all_outputs = []
+        for i, h in enumerate(layer_rep):
+            pooled_h = self.pool(h, batch.batch)
+            all_outputs.append(pooled_h)
+            score_over_layer += self.final_drop(self.linear_prediction[i](pooled_h))
+
+        return score_over_layer, all_outputs[1:]
+
+
+@register_model("gcc")
+class GraphEncoder(BaseModel):
+    """
+    MPNN from
+    `Neural Message Passing for Quantum Chemistry <https://arxiv.org/abs/1704.01212>`__
+    Parameters
+    ----------
+    node_input_dim : int
+        Dimension of input node feature, default to be 15.
+    edge_input_dim : int
+        Dimension of input edge feature, default to be 15.
+    output_dim : int
+        Dimension of prediction, default to be 12.
+    node_hidden_dim : int
+        Dimension of node feature in hidden layers, default to be 64.
+    edge_hidden_dim : int
+        Dimension of edge feature in hidden layers, default to be 128.
+    num_step_message_passing : int
+        Number of message passing steps, default to be 6.
+    num_step_set2set : int
+        Number of set2set steps
+    num_layer_set2set : int
+        Number of set2set layers
+    """
+
+    @staticmethod
+    def add_args(parser):
+        parser.add_argument("--hidden-size", type=int, default=64)
+        parser.add_argument("--positional-embedding-size", type=int, default=32)
+        parser.add_argument("--degree-embedding-size", type=int, default=16)
+        parser.add_argument("--max-node-freq", type=int, default=16)
+        parser.add_argument("--max-edge-freq", type=int, default=16)
+        parser.add_argument("--max-degree", type=int, default=512)
+        parser.add_argument("--freq-embedding-size", type=int, default=16)
+        parser.add_argument("--num-layers", type=int, default=2)
+        parser.add_argument("--num-heads", type=int, default=2)
+        parser.add_argument("--output-size", type=int, default=32)
+
+    @classmethod
+    def build_model_from_args(cls, args):
+        return cls(
+            positional_embedding_size=args.positional_embedding_size,
+            max_node_freq=args.max_node_freq,
+            max_edge_freq=args.max_edge_freq,
+            max_degree=args.max_degree,
+            num_layers=args.num_layers,
+            num_heads=args.num_heads,
+            degree_embedding_size=args.degree_embedding_size,
+            node_hidden_dim=args.hidden_size,
+            output_dim=args.output_size,
+        )
+
+    def __init__(
+        self,
+        positional_embedding_size=32,
+        max_node_freq=8,
+        max_edge_freq=8,
+        max_degree=128,
+        freq_embedding_size=32,
+        degree_embedding_size=32,
+        output_dim=32,
+        node_hidden_dim=32,
+        edge_hidden_dim=32,
+        num_layers=6,
+        num_heads=4,
+        num_step_set2set=6,
+        num_layer_set2set=3,
+        norm=False,
+        gnn_model="gin",
+        degree_input=False,
+    ):
+        super(GraphEncoder, self).__init__()
+
+        if degree_input:
+            node_input_dim = positional_embedding_size + degree_embedding_size + 1
+        else:
+            node_input_dim = positional_embedding_size + 1
+        # node_input_dim = (
+        #     positional_embedding_size + freq_embedding_size + degree_embedding_size + 3
+        # )
+        # edge_input_dim = freq_embedding_size + 1
+        if gnn_model == "gat":
+            self.gnn = GATModel(
+                in_feats=node_input_dim,
+                hidden_size=node_hidden_dim,
+                num_layers=num_layers,
+                nhead=num_heads,
+                dropout=0.0,
+            )
+        elif gnn_model == "gin":
+            self.gnn = GINModel(
+                num_layers=num_layers,
+                num_mlp_layers=2,
+                in_feats=node_input_dim,
+                hidden_dim=node_hidden_dim,
+                out_feats=output_dim,
+                final_dropout=0.5,
+                train_eps=False,
+                pooling="sum",
+                # neighbor_pooling_type="sum",
+                # use_selayer=False,
+            )
+        self.gnn_model = gnn_model
+
+        self.max_node_freq = max_node_freq
+        self.max_edge_freq = max_edge_freq
+        self.max_degree = max_degree
+        self.degree_input = degree_input
+
+        # self.node_freq_embedding = nn.Embedding(
+        #     num_embeddings=max_node_freq + 1, embedding_dim=freq_embedding_size
+        # )
+        if degree_input:
+            self.degree_embedding = nn.Embedding(num_embeddings=max_degree + 1, embedding_dim=degree_embedding_size)
+
+        # self.edge_freq_embedding = nn.Embedding(
+        #     num_embeddings=max_edge_freq + 1, embedding_dim=freq_embedding_size
+        # )
+
+        self.set2set = Set2Set(node_hidden_dim, num_step_set2set, num_layer_set2set)
+        if gnn_model != "gin":
+            self.lin_readout = nn.Sequential(
+                nn.Linear(2 * node_hidden_dim, node_hidden_dim),
+                nn.ReLU(),
+                nn.Linear(node_hidden_dim, output_dim),
+            )
+        else:
+            self.lin_readout = None
+        self.norm = norm
+
+    def forward(self, g, return_all_outputs=False):
+        """Predict molecule labels
+        Parameters
+        ----------
+        g : Graph
+        n_feat : tensor of dtype float32 and shape (B1, D1)
+            Node features. B1 for number of nodes and D1 for
+            the node feature size.
+        e_feat : tensor of dtype float32 and shape (B2, D2)
+            Edge features. B2 for number of edges and D2 for
+            the edge feature size.
+        Returns
+        -------
+        res : Predicted labels
+        """
+
+        # nfreq = g.ndata["nfreq"]
+        device = self.device
+        pos_undirected = g.pos_undirected
+        seed_emb = g.seed.unsqueeze(1).float()
+        if not torch.is_tensor(seed_emb):
+            seed_emb = torch.Tensor(seed_emb)
+
+        if self.degree_input:
+            degrees = g.degrees()
+            if device != torch.device("cpu"):
+                degrees = degrees.cuda(device)
+
+            deg_emb = self.degree_embedding(degrees.clamp(0, self.max_degree))
+
+            n_feat = torch.cat((pos_undirected, deg_emb, seed_emb), dim=-1)
+        else:
+            n_feat = torch.cat(
+                (
+                    pos_undirected,
+                    # self.node_freq_embedding(nfreq.clamp(0, self.max_node_freq)),
+                    # self.degree_embedding(degrees.clamp(0, self.max_degree)),
+                    seed_emb,
+                    # nfreq.unsqueeze(1).float() / self.max_node_freq,
+                    # degrees.unsqueeze(1).float() / self.max_degree,
+                ),
+                dim=-1,
+            )
+
+        if self.gnn_model == "gin":
+            x, all_outputs = self.gnn(g, n_feat)
+        else:
+            x, all_outputs = self.gnn(g, n_feat), None
+            x = self.set2set(g, x)
+            x = self.lin_readout(x)
+        if self.norm:
+            x = F.normalize(x, p=2, dim=-1, eps=1e-5)
+        if return_all_outputs:
+            return x, all_outputs
+        else:
+            return x
+
+
+# --------------------------------------------------------------------
+# --------------------------------------------------------------------
+
+
+def warmup_linear(x, warmup=0.002):
+    """Specifies a triangular learning rate schedule where peak is reached at `warmup`*`t_total`-th (as provided to BertAdam) training step.
+    After `t_total`-th training step, learning rate is zero."""
+    if x < warmup:
+        return x / warmup
+    return max((x - 1.0) / (warmup - 1.0), 0)
diff --git a/cogdl/models/nn/gcn.py b/cogdl/models/nn/gcn.py
index dba49231..7fe02245 100644
--- a/cogdl/models/nn/gcn.py
+++ b/cogdl/models/nn/gcn.py
@@ -28,6 +28,7 @@ def add_args(parser):
         parser.add_argument("--residual", action="store_true")
         parser.add_argument("--norm", type=str, default=None)
         parser.add_argument("--activation", type=str, default="relu")
+        parser.add_argument("--actnn", action="store_true")
         # fmt: on
 
     @classmethod
@@ -41,7 +42,7 @@ def build_model_from_args(cls, args):
             args.activation,
             args.residual,
             args.norm,
-            args.actnn,
+            args.actnn if hasattr(args, "actnn") else False,
         )
 
     def __init__(
diff --git a/cogdl/models/nn/gcnmix.py b/cogdl/models/nn/gcnmix.py
index 0b619655..9fd007be 100644
--- a/cogdl/models/nn/gcnmix.py
+++ b/cogdl/models/nn/gcnmix.py
@@ -61,7 +61,33 @@ def forward_aux(self, x):
         return self.weight(x)
 
 
+@register_model("gcnmix")
 class BaseGNNMix(BaseModel):
+    @staticmethod
+    def add_args(parser):
+        parser.add_argument("--dropout", type=float, default=0.5)
+        parser.add_argument("--hidden-size", type=int, default=64)
+        parser.add_argument("--alpha", type=float, default=1.0)
+        parser.add_argument("--k", type=int, default=10)
+        parser.add_argument("--temperature", type=float, default=0.1)
+        # parser.add_argument("--rampup-starts", type=int, default=500)
+        # parser.add_argument("--rampup_ends", type=int, default=1000)
+        # parser.add_argument("--mixup-consistency", type=float, default=10.0)
+        # parser.add_argument("--ema-decay", type=float, default=0.999)
+        # parser.add_argument("--tau", type=float, default=1.0)
+
+    @classmethod
+    def build_model_from_args(cls, args):
+        return cls(
+            in_feat=args.num_features,
+            hidden_size=args.hidden_size,
+            num_classes=args.num_classes,
+            k=args.k,
+            temperature=args.temperature,
+            alpha=args.alpha,
+            dropout=args.dropout,
+        )
+
     def __init__(self, in_feat, hidden_size, num_classes, k, temperature, alpha, dropout):
         super(BaseGNNMix, self).__init__()
         self.dropout = dropout
@@ -147,86 +173,86 @@ def predict_noise(self, data, tau=1):
         return out
 
 
-@register_model("gcnmix")
-class GCNMix(BaseModel):
-    @staticmethod
-    def add_args(parser):
-        parser.add_argument("--dropout", type=float, default=0.5)
-        parser.add_argument("--hidden-size", type=int, default=64)
-        parser.add_argument("--alpha", type=float, default=1.0)
-        parser.add_argument("--k", type=int, default=10)
-        parser.add_argument("--temperature", type=float, default=0.1)
-        parser.add_argument("--rampup-starts", type=int, default=500)
-        parser.add_argument("--rampup_ends", type=int, default=1000)
-        parser.add_argument("--mixup-consistency", type=float, default=10.0)
-        parser.add_argument("--ema-decay", type=float, default=0.999)
-        parser.add_argument("--tau", type=float, default=1.0)
-
-    @classmethod
-    def build_model_from_args(cls, args):
-        return cls(
-            in_feat=args.num_features,
-            hidden_size=args.hidden_size,
-            num_classes=args.num_classes,
-            k=args.k,
-            temperature=args.temperature,
-            alpha=args.alpha,
-            rampup_starts=args.rampup_starts,
-            rampup_ends=args.rampup_ends,
-            final_consistency_weight=args.mixup_consistency,
-            ema_decay=args.ema_decay,
-            dropout=args.dropout,
-        )
-
-    def __init__(
-        self,
-        in_feat,
-        hidden_size,
-        num_classes,
-        k,
-        temperature,
-        alpha,
-        rampup_starts,
-        rampup_ends,
-        final_consistency_weight,
-        ema_decay,
-        dropout,
-    ):
-        super(GCNMix, self).__init__()
-        self.final_consistency_weight = final_consistency_weight
-        self.rampup_starts = rampup_starts
-        self.rampup_ends = rampup_ends
-        self.ema_decay = ema_decay
-
-        self.base_gnn = BaseGNNMix(in_feat, hidden_size, num_classes, k, temperature, alpha, dropout)
-        self.ema_gnn = BaseGNNMix(in_feat, hidden_size, num_classes, k, temperature, alpha, dropout)
-        for param in self.ema_gnn.parameters():
-            param.detach_()
-
-        self.epoch = 0
-
-    def forward(self, graph):
-        return self.base_gnn.forward(graph)
-
-    def forward_ema(self, graph):
-        return self.ema_gnn(graph)
-
-    def node_classification_loss(self, data):
-        opt = {
-            "epoch": self.epoch,
-            "final_consistency_weight": self.final_consistency_weight,
-            "rampup_starts": self.rampup_starts,
-            "rampup_ends": self.rampup_ends,
-        }
-        self.base_gnn.train()
-        loss_n = self.base_gnn.loss(data, opt)
-
-        alpha = min(1 - 1 / (self.epoch + 1), self.ema_decay)
-        for ema_param, param in zip(self.ema_gnn.parameters(), self.base_gnn.parameters()):
-            ema_param.data.mul_(alpha).add_((1 - alpha) * param.data)
-        self.epoch += 1
-        return loss_n
-
-    def predict(self, data):
-        prediction = self.forward_ema(data)
-        return prediction
+# @register_model("gcnmix")
+# class GCNMix(BaseModel):
+#     @staticmethod
+#     def add_args(parser):
+#         parser.add_argument("--dropout", type=float, default=0.5)
+#         parser.add_argument("--hidden-size", type=int, default=64)
+#         parser.add_argument("--alpha", type=float, default=1.0)
+#         parser.add_argument("--k", type=int, default=10)
+#         parser.add_argument("--temperature", type=float, default=0.1)
+#         parser.add_argument("--rampup-starts", type=int, default=500)
+#         parser.add_argument("--rampup_ends", type=int, default=1000)
+#         parser.add_argument("--mixup-consistency", type=float, default=10.0)
+#         parser.add_argument("--ema-decay", type=float, default=0.999)
+#         parser.add_argument("--tau", type=float, default=1.0)
+#
+#     @classmethod
+#     def build_model_from_args(cls, args):
+#         return cls(
+#             in_feat=args.num_features,
+#             hidden_size=args.hidden_size,
+#             num_classes=args.num_classes,
+#             k=args.k,
+#             temperature=args.temperature,
+#             alpha=args.alpha,
+#             rampup_starts=args.rampup_starts,
+#             rampup_ends=args.rampup_ends,
+#             final_consistency_weight=args.mixup_consistency,
+#             ema_decay=args.ema_decay,
+#             dropout=args.dropout,
+#         )
+#
+#     def __init__(
+#         self,
+#         in_feat,
+#         hidden_size,
+#         num_classes,
+#         k,
+#         temperature,
+#         alpha,
+#         rampup_starts,
+#         rampup_ends,
+#         final_consistency_weight,
+#         ema_decay,
+#         dropout,
+#     ):
+#         super(GCNMix, self).__init__()
+#         self.final_consistency_weight = final_consistency_weight
+#         self.rampup_starts = rampup_starts
+#         self.rampup_ends = rampup_ends
+#         self.ema_decay = ema_decay
+#
+#         self.base_gnn = BaseGNNMix(in_feat, hidden_size, num_classes, k, temperature, alpha, dropout)
+#         self.ema_gnn = BaseGNNMix(in_feat, hidden_size, num_classes, k, temperature, alpha, dropout)
+#         for param in self.ema_gnn.parameters():
+#             param.detach_()
+#
+#         self.epoch = 0
+#
+#     def forward(self, graph):
+#         return self.base_gnn.forward(graph)
+#
+#     def forward_ema(self, graph):
+#         return self.ema_gnn(graph)
+#
+#     def node_classification_loss(self, data):
+#         opt = {
+#             "epoch": self.epoch,
+#             "final_consistency_weight": self.final_consistency_weight,
+#             "rampup_starts": self.rampup_starts,
+#             "rampup_ends": self.rampup_ends,
+#         }
+#         self.base_gnn.train()
+#         loss_n = self.base_gnn.loss(data, opt)
+#
+#         alpha = min(1 - 1 / (self.epoch + 1), self.ema_decay)
+#         for ema_param, param in zip(self.ema_gnn.parameters(), self.base_gnn.parameters()):
+#             ema_param.data.mul_(alpha).add_((1 - alpha) * param.data)
+#         self.epoch += 1
+#         return loss_n
+#
+#     def predict(self, data):
+#         prediction = self.forward_ema(data)
+#         return prediction
diff --git a/cogdl/models/nn/gin.py b/cogdl/models/nn/gin.py
index 8100d8a0..675cc894 100644
--- a/cogdl/models/nn/gin.py
+++ b/cogdl/models/nn/gin.py
@@ -41,10 +41,6 @@ def add_args(parser):
         parser.add_argument("--dropout", type=float, default=0.5)
         parser.add_argument("--train-epsilon", dest="train_epsilon", action="store_false")
         parser.add_argument("--pooling", type=str, default="sum")
-        parser.add_argument("--batch-size", type=int, default=128)
-        parser.add_argument("--lr", type=float, default=0.001)
-        parser.add_argument("--train-ratio", type=float, default=0.7)
-        parser.add_argument("--test-ratio", type=float, default=0.1)
 
     @classmethod
     def build_model_from_args(cls, args):
diff --git a/cogdl/models/nn/grand.py b/cogdl/models/nn/grand.py
index 1d6b6add..c40575a8 100644
--- a/cogdl/models/nn/grand.py
+++ b/cogdl/models/nn/grand.py
@@ -51,11 +51,8 @@ def add_args(parser):
         parser.add_argument("--input-dropout", type=float, default=0.5)
         parser.add_argument("--bn", type=bool, default=False)
         parser.add_argument("--dropnode-rate", type=float, default=0.5)
-        parser.add_argument('--order', type=int, default=5)
-        parser.add_argument('--tem', type=float, default=0.5)
-        parser.add_argument('--lam', type=float, default=0.5)
-        parser.add_argument('--sample', type=int, default=2)    
-        parser.add_argument('--alpha', type=float, default=0.2)        
+        parser.add_argument("--order", type=int, default=5)
+        parser.add_argument("--alpha", type=float, default=0.2)
 
         # fmt: on
 
@@ -69,10 +66,7 @@ def build_model_from_args(cls, args):
             args.hidden_dropout,
             args.bn,
             args.dropnode_rate,
-            args.tem,
-            args.lam,
             args.order,
-            args.sample,
             args.alpha,
         )
 
@@ -85,10 +79,7 @@ def __init__(
         hidden_droprate,
         use_bn,
         dropnode_rate,
-        tem,
-        lam,
         order,
-        sample,
         alpha,
     ):
         super(Grand, self).__init__()
@@ -99,14 +90,11 @@ def __init__(
         self.bn1 = nn.BatchNorm1d(nfeat)
         self.bn2 = nn.BatchNorm1d(nhid)
         self.use_bn = use_bn
-        self.tem = tem
-        self.lam = lam
         self.order = order
         self.dropnode_rate = dropnode_rate
-        self.sample = sample
         self.alpha = alpha
 
-    def dropNode(self, x):
+    def drop_node(self, x):
         n = x.shape[0]
         drop_rates = torch.ones(n) * self.dropnode_rate
         if self.training:
@@ -118,7 +106,7 @@ def dropNode(self, x):
         return x
 
     def rand_prop(self, graph, x):
-        x = self.dropNode(x)
+        x = self.drop_node(x)
 
         y = x
         for i in range(self.order):
@@ -126,21 +114,6 @@ def rand_prop(self, graph, x):
             y.add_(x)
         return y.div_(self.order + 1.0).detach_()
 
-    def consis_loss(self, logps, train_mask):
-        temp = self.tem
-        ps = [torch.exp(p)[~train_mask] for p in logps]
-        sum_p = 0.0
-        for p in ps:
-            sum_p = sum_p + p
-        avg_p = sum_p / len(ps)
-        sharp_p = (torch.pow(avg_p, 1.0 / temp) / torch.sum(torch.pow(avg_p, 1.0 / temp), dim=1, keepdim=True)).detach()
-        loss = 0.0
-        for p in ps:
-            loss += torch.mean((p - sharp_p).pow(2).sum(1))
-        loss = loss / len(ps)
-
-        return self.lam * loss
-
     def normalize_x(self, x):
         row_sum = x.sum(1)
         row_inv = row_sum.pow_(-1)
@@ -163,22 +136,5 @@ def forward(self, graph):
         x = self.layer2(x)
         return x
 
-    def node_classification_loss(self, graph):
-        output_list = []
-        for i in range(self.sample):
-            output_list.append(self.forward(graph))
-        loss_train = 0.0
-        for output in output_list:
-            loss_train += self.loss_fn(output[graph.train_mask], graph.y[graph.train_mask])
-        loss_train = loss_train / self.sample
-
-        if len(graph.y.shape) > 1:
-            output_list = [torch.sigmoid(x) for x in output_list]
-        else:
-            output_list = [F.log_softmax(x, dim=-1) for x in output_list]
-        loss_consis = self.consis_loss(output_list, graph.train_mask)
-
-        return loss_train + loss_consis
-
     def predict(self, data):
         return self.forward(data)
diff --git a/cogdl/models/nn/graphsage.py b/cogdl/models/nn/graphsage.py
index d1cc737c..bf7d3ac4 100644
--- a/cogdl/models/nn/graphsage.py
+++ b/cogdl/models/nn/graphsage.py
@@ -1,4 +1,3 @@
-from typing import Any
 import random
 
 import torch
@@ -45,7 +44,7 @@ def add_args(parser):
         parser.add_argument("--num-layers", type=int, default=2)
         parser.add_argument("--sample-size", type=int, nargs='+', default=[10, 10])
         parser.add_argument("--dropout", type=float, default=0.5)
-        parser.add_argument("--batch-size", type=int, default=128)
+        # parser.add_argument("--batch-size", type=int, default=128)
         parser.add_argument("--aggr", type=str, default="mean")
         # fmt: on
 
diff --git a/cogdl/models/nn/han.py b/cogdl/models/nn/han.py
index 48dcc363..72473261 100644
--- a/cogdl/models/nn/han.py
+++ b/cogdl/models/nn/han.py
@@ -50,20 +50,10 @@ def __init__(self, num_edge, w_in, w_out, num_class, num_nodes, num_layers):
         self.cross_entropy_loss = nn.CrossEntropyLoss()
         self.linear = nn.Linear(self.w_out, self.num_class)
 
-    def forward(self, graph, target_x, target):
+    def forward(self, graph):
         X = graph.x
         for i in range(self.num_layers):
             X = self.layers[i](graph, X)
 
-        y = self.linear(X[target_x])
-        loss = self.cross_entropy_loss(y, target)
-        return loss, y
-
-    def loss(self, data):
-        loss, y = self.forward(data, data.train_node, data.train_target)
-        return loss
-
-    def evaluate(self, data, nodes, targets):
-        loss, y = self.forward(data, nodes, targets)
-        f1 = accuracy(y, targets)
-        return loss.item(), f1
+        out = self.linear(X)
+        return out
diff --git a/cogdl/models/nn/infograph.py b/cogdl/models/nn/infograph.py
index d3941d1c..390c40fe 100644
--- a/cogdl/models/nn/infograph.py
+++ b/cogdl/models/nn/infograph.py
@@ -131,6 +131,7 @@ def __init__(self, in_feats, hidden_dim, out_feats, num_layers=3, sup=False):
         self.sup = sup
         self.emb_dim = hidden_dim
         self.out_feats = out_feats
+        self.num_layers = num_layers
 
         self.sem_fc1 = nn.Linear(num_layers * hidden_dim, hidden_dim)
         self.sem_fc2 = nn.Linear(hidden_dim, out_feats)
diff --git a/cogdl/models/nn/mvgrl.py b/cogdl/models/nn/mvgrl.py
index 8443e840..c5174dba 100644
--- a/cogdl/models/nn/mvgrl.py
+++ b/cogdl/models/nn/mvgrl.py
@@ -129,6 +129,7 @@ def augment(self, graph):
         return adj, diff
 
     def preprocess(self, graph):
+        print("MVGRL preprocessing...")
         graph.add_remaining_self_loops()
         graph.sym_norm()
 
@@ -147,8 +148,12 @@ def preprocess(self, graph):
         self.cache["diff"] = graphs[1]
         self.cache["adj"] = graphs[0]
         self.device = next(self.gcn1.parameters()).device
+        print("Preprocessing Done...")
 
     def forward(self, graph):
+        if not self.training:
+            return self.embed(graph)
+
         x = graph.x
         if self.cache is None or "diff" not in self.cache:
             self.preprocess(graph)
@@ -182,9 +187,6 @@ def loss(self, data):
         loss = self.loss_f(logits, lbl)
         return loss
 
-    def self_supervised_loss(self, data):
-        return self.loss(data)
-
     def embed(self, data, msk=None):
         adj = self.cache["adj"].to(self.device)
         diff = self.cache["diff"].to(self.device)
diff --git a/cogdl/models/nn/patchy_san.py b/cogdl/models/nn/patchy_san.py
index 0159960c..5e547e8b 100644
--- a/cogdl/models/nn/patchy_san.py
+++ b/cogdl/models/nn/patchy_san.py
@@ -26,48 +26,44 @@ class PatchySAN(BaseModel):
 
     @staticmethod
     def add_args(parser):
-        parser.add_argument("--batch-size", type=int, default=20)
-        parser.add_argument("--sample", default=30, type=int, help="Number of chosen vertexes")
-        parser.add_argument("--stride", default=1, type=int, help="Stride of chosen vertexes")
-        parser.add_argument("--neighbor", default=10, type=int, help="Number of neighbor in constructing features")
+        parser.add_argument("--num-sample", default=30, type=int, help="Number of chosen vertexes")
+        # parser.add_argument("--stride", default=1, type=int, help="Stride of chosen vertexes")
+        parser.add_argument("--num-neighbor", default=10, type=int, help="Number of neighbor in constructing features")
         parser.add_argument("--iteration", default=5, type=int, help="Number of iteration")
-        parser.add_argument("--train-ratio", type=float, default=0.7)
-        parser.add_argument("--test-ratio", type=float, default=0.1)
 
     @classmethod
     def build_model_from_args(cls, args):
         return cls(
-            args.batch_size,
             args.num_features,
             args.num_classes,
             args.sample,
-            args.stride,
+            # args.stride,
             args.neighbor,
             args.iteration,
         )
 
     @classmethod
-    def split_dataset(self, dataset, args):
-        # process each graph and add it into Data() as attribute tx
+    def split_dataset(cls, dataset, args):
+        # process each graph and add it into Data() as attribute x
         for i, data in enumerate(dataset):
             new_feature = get_single_feature(
                 dataset[i], args.num_features, args.num_classes, args.sample, args.neighbor, args.stride
             )
-            dataset[i].tx = torch.from_numpy(new_feature)
+            dataset[i].x = torch.from_numpy(new_feature)
 
         return split_dataset_general(dataset, args)
 
-    def __init__(self, batch_size, num_features, num_classes, num_sample, stride, num_neighbor, iteration):
+    # def __init__(self, batch_size, num_features, num_classes, num_sample, stride, num_neighbor, iteration):
+    def __init__(self, num_features, num_classes, num_sample, num_neighbor, iteration):
         super(PatchySAN, self).__init__()
 
-        self.batch_size = batch_size
         self.num_features = num_features
         self.num_classes = num_classes
         self.num_sample = num_sample
-        self.stride = stride
         self.num_neighbor = num_neighbor
         self.iteration = iteration
 
+        # self.build_model(self.num_features, self.num_sample, self.num_neighbor, self.num_classes)
         self.build_model(self.num_features, self.num_sample, self.num_neighbor, self.num_classes)
 
     def build_model(self, num_channel, num_sample, num_neighbor, num_class):
@@ -94,7 +90,7 @@ def build_model(self, num_channel, num_sample, num_neighbor, num_class):
         )
 
     def forward(self, batch):
-        logits = self.nn(batch.tx)
+        logits = self.nn(batch.x)
         return logits
 
 
diff --git a/cogdl/models/nn/pprgo.py b/cogdl/models/nn/pprgo.py
index 63629716..ac05c99a 100644
--- a/cogdl/models/nn/pprgo.py
+++ b/cogdl/models/nn/pprgo.py
@@ -16,15 +16,7 @@ def add_args(parser):
         parser.add_argument("--dropout", type=float, default=0.1)
         parser.add_argument("--activation", type=str, default="relu")
         parser.add_argument("--nprop-inference", type=int, default=2)
-
         parser.add_argument("--alpha", type=float, default=0.5)
-        parser.add_argument("--k", type=int, default=32)
-        parser.add_argument("--norm", type=str, default="sym")
-        parser.add_argument("--eps", type=float, default=1e-4)
-
-        parser.add_argument("--eval-step", type=int, default=4)
-        parser.add_argument("--batch-size", type=int, default=512)
-        parser.add_argument("--test-batch-size", type=int, default=10000)
 
     @classmethod
     def build_model_from_args(cls, args):
@@ -37,11 +29,15 @@ def build_model_from_args(cls, args):
             dropout=args.dropout,
             activation=args.activation,
             nprop=args.nprop_inference,
+            norm=args.norm if hasattr(args, "norm") else "sym",
         )
 
-    def __init__(self, in_feats, hidden_size, out_feats, num_layers, alpha, dropout, activation="relu", nprop=2):
+    def __init__(
+        self, in_feats, hidden_size, out_feats, num_layers, alpha, dropout, activation="relu", nprop=2, norm="sym"
+    ):
         super(PPRGo, self).__init__()
         self.alpha = alpha
+        self.norm = norm
         self.nprop = nprop
         self.fc = PPRGoLayer(in_feats, hidden_size, out_feats, num_layers, dropout, activation)
 
@@ -58,7 +54,7 @@ def node_classification_loss(self, x, targets, ppr_scores, y):
         loss = self.loss_fn(pred, y)
         return loss
 
-    def predict(self, graph, batch_size, norm):
+    def predict(self, graph, batch_size=10000):
         device = next(self.fc.parameters()).device
         x = graph.x
         num_nodes = x.shape[0]
@@ -71,9 +67,9 @@ def predict(self, graph, batch_size, norm):
         pred_logits = torch.cat(pred_logits, dim=0)
 
         with graph.local_graph():
-            if norm == "sym":
+            if self.norm == "sym":
                 graph.sym_norm()
-            elif norm == "row":
+            elif self.norm == "row":
                 graph.row_norm()
             else:
                 raise NotImplementedError
diff --git a/cogdl/models/nn/pyg_gpt_gnn.py b/cogdl/models/nn/pyg_gpt_gnn.py
deleted file mode 100644
index 68e1f5cc..00000000
--- a/cogdl/models/nn/pyg_gpt_gnn.py
+++ /dev/null
@@ -1,209 +0,0 @@
-from typing import Any, Union, Type, Optional
-
-from cogdl.models import register_model
-from cogdl.models.supervised_model import (
-    SupervisedHomogeneousNodeClassificationModel,
-    SupervisedHeterogeneousNodeClassificationModel,
-)
-
-from cogdl.trainers.gpt_gnn_trainer import (
-    GPT_GNNHomogeneousTrainer,
-    GPT_GNNHeterogeneousTrainer,
-)
-
-
-#
-# @register_model("gpt_gnn")
-# class GPT_GNN(BaseModel):
-#     def __init__(
-#         self,
-#         in_dim,
-#         n_hid,
-#         num_types,
-#         num_relations,
-#         n_heads,
-#         n_layers,
-#         dropout=0.2,
-#         conv_name="hgt",
-#         prev_norm=False,
-#         last_norm=False,
-#         use_RTE=True,
-#     ):
-#         super(GPT_GNN, self).__init__()
-#         self.gcs = nn.ModuleList()
-#         self.num_types = num_types
-#         self.in_dim = in_dim
-#         self.n_hid = n_hid
-#         self.adapt_ws = nn.ModuleList()
-#         self.drop = nn.Dropout(dropout)
-#         for t in range(num_types):
-#             self.adapt_ws.append(nn.Linear(in_dim, n_hid))
-#         for l in range(n_layers - 1):
-#             self.gcs.append(
-#                 GeneralConv(
-#                     conv_name,
-#                     n_hid,
-#                     n_hid,
-#                     num_types,
-#                     num_relations,
-#                     n_heads,
-#                     dropout,
-#                     use_norm=prev_norm,
-#                     use_RTE=use_RTE,
-#                 )
-#             )
-#         self.gcs.append(
-#             GeneralConv(
-#                 conv_name,
-#                 n_hid,
-#                 n_hid,
-#                 num_types,
-#                 num_relations,
-#                 n_heads,
-#                 dropout,
-#                 use_norm=last_norm,
-#                 use_RTE=use_RTE,
-#             )
-#         )
-#
-#     def forward(self, node_feature, node_type, edge_time, edge_index, edge_type):
-#         res = torch.zeros(node_feature.size(0), self.n_hid).to(node_feature.device)
-#         for t_id in range(self.num_types):
-#             idx = node_type == int(t_id)
-#             if idx.sum() == 0:
-#                 continue
-#             res[idx] = torch.tanh(self.adapt_ws[t_id](node_feature[idx]))
-#         meta_xs = self.drop(res)
-#         del res
-#         for gc in self.gcs:
-#             meta_xs = gc(meta_xs, node_type, edge_index, edge_type, edge_time)
-#         return meta_xs
-
-
-@register_model("gpt_gnn")
-class GPT_GNN(
-    SupervisedHomogeneousNodeClassificationModel,
-    SupervisedHeterogeneousNodeClassificationModel,
-):
-    @staticmethod
-    def add_args(parser):
-        """Add task-specific arguments to the parser."""
-        # fmt: off
-        """
-            Dataset arguments
-        """
-        parser.add_argument(
-            "--use_pretrain", help="Whether to use pre-trained model", action="store_true"
-        )
-        parser.add_argument(
-            "--pretrain_model_dir",
-            type=str,
-            default="/datadrive/models/gpt_all_cs",
-            help="The address for pretrained model.",
-        )
-        # parser.add_argument(
-        #     "--model_dir",
-        #     type=str,
-        #     default="/datadrive/models/gpt_all_reddit",
-        #     help="The address for storing the models and optimization results.",
-        # )
-        parser.add_argument(
-            "--task_name",
-            type=str,
-            default="reddit",
-            help="The name of the stored models and optimization results.",
-        )
-        parser.add_argument(
-            "--sample_depth", type=int, default=6, help="How many numbers to sample the graph"
-        )
-        parser.add_argument(
-            "--sample_width",
-            type=int,
-            default=128,
-            help="How many nodes to be sampled per layer per type",
-        )
-        """
-           Model arguments 
-        """
-        parser.add_argument(
-            "--conv_name",
-            type=str,
-            default="hgt",
-            choices=["hgt", "gcn", "gat", "rgcn", "han", "hetgnn"],
-            help="The name of GNN filter. By default is Heterogeneous Graph Transformer (hgt)",
-        )
-        parser.add_argument("--n_hid", type=int, default=400, help="Number of hidden dimension")
-        parser.add_argument("--n_heads", type=int, default=8, help="Number of attention head")
-        parser.add_argument("--n_layers", type=int, default=3, help="Number of GNN layers")
-        parser.add_argument(
-            "--prev_norm",
-            help="Whether to add layer-norm on the previous layers",
-            action="store_true",
-        )
-        parser.add_argument(
-            "--last_norm",
-            help="Whether to add layer-norm on the last layers",
-            action="store_true",
-        )
-        parser.add_argument("--dropout", type=int, default=0.2, help="Dropout ratio")
-
-        """
-            Optimization arguments
-        """
-        parser.add_argument(
-            "--optimizer",
-            type=str,
-            default="adamw",
-            choices=["adamw", "adam", "sgd", "adagrad"],
-            help="optimizer to use.",
-        )
-        parser.add_argument(
-            "--scheduler",
-            type=str,
-            default="cosine",
-            help="Name of learning rate scheduler.",
-            choices=["cycle", "cosine"],
-        )
-        parser.add_argument(
-            "--data_percentage",
-            type=int,
-            default=0.1,
-            help="Percentage of training and validation data to use",
-        )
-        parser.add_argument("--n_epoch", type=int, default=50, help="Number of epoch to run")
-        parser.add_argument(
-            "--n_pool", type=int, default=8, help="Number of process to sample subgraph"
-        )
-        parser.add_argument(
-            "--n_batch",
-            type=int,
-            default=10,
-            help="Number of batch (sampled graphs) for each epoch",
-        )
-        parser.add_argument(
-            "--batch_size", type=int, default=64, help="Number of output nodes for training"
-        )
-        parser.add_argument("--clip", type=int, default=0.5, help="Gradient Norm Clipping")
-        # fmt: on
-
-    @classmethod
-    def build_model_from_args(cls, args):
-        return GPT_GNN()
-
-    def loss(self, data: Any) -> Any:
-        pass
-
-    def predict(self, data: Any) -> Any:
-        pass
-
-    def evaluate(self, data: Any, nodes: Any, targets: Any) -> Any:
-        pass
-
-    @staticmethod
-    def get_trainer(args) -> Optional[Type[Union[GPT_GNNHomogeneousTrainer, GPT_GNNHeterogeneousTrainer]]]:
-        # if taskType == NodeClassification:
-        return GPT_GNNHomogeneousTrainer
-        # elif taskType == HeterogeneousNodeClassification:
-        #     return GPT_GNNHeterogeneousTrainer
-        # else:
-        #     return None
diff --git a/cogdl/models/nn/pyg_gtn.py b/cogdl/models/nn/pyg_gtn.py
index b76808d7..982bef10 100644
--- a/cogdl/models/nn/pyg_gtn.py
+++ b/cogdl/models/nn/pyg_gtn.py
@@ -157,7 +157,7 @@ def norm(self, edge_index, num_nodes, edge_weight, improved=False, dtype=None):
 
         return deg_inv_sqrt[row], deg_inv_sqrt[col]
 
-    def forward(self, graph, target_x, target):
+    def forward(self, graph):
         A = graph.adj
         X = graph.x
         Ws = []
@@ -183,16 +183,5 @@ def forward(self, graph, target_x, target):
                     X_ = torch.cat((X_, F.relu(self.gcn(graph, X))), dim=1)
         X_ = self.linear1(X_)
         X_ = F.relu(X_)
-        # X_ = F.dropout(X_, p=0.5)
-        y = self.linear2(X_[target_x])
-        loss = self.cross_entropy_loss(y, target)
-        return loss, y, Ws
-
-    def loss(self, data):
-        loss, y, _ = self.forward(data, data.train_node, data.train_target)
-        return loss
-
-    def evaluate(self, data, nodes, targets):
-        loss, y, _ = self.forward(data, nodes, targets)
-        f1 = accuracy(y, targets)
-        return loss.item(), f1
+        out = self.linear2(X_)
+        return out
diff --git a/cogdl/models/nn/rgcn.py b/cogdl/models/nn/rgcn.py
index 6d8441be..be371403 100644
--- a/cogdl/models/nn/rgcn.py
+++ b/cogdl/models/nn/rgcn.py
@@ -2,7 +2,7 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
-from cogdl.utils.link_prediction_utils import GNNLinkPredict, cal_mrr, sampling_edge_uniform
+from cogdl.utils.link_prediction_utils import GNNLinkPredict, sampling_edge_uniform
 from cogdl.layers import RGCNLayer
 from .. import register_model, BaseModel
 
@@ -84,7 +84,7 @@ def __init__(
         self_dropout=0.0,
     ):
         BaseModel.__init__(self)
-        GNNLinkPredict.__init__(self, "distmult", hidden_size)
+        GNNLinkPredict.__init__(self)
         self.penalty = penalty
         self.num_nodes = num_entities
         self.num_rels = num_rels
@@ -120,15 +120,9 @@ def forward(self, graph):
         # output = self.model(x, reindexed_indices, graph.edge_type)
         return output
 
-    def loss(self, graph, split="train"):
-        if split == "train":
-            mask = graph.train_mask
-        elif split == "val":
-            mask = graph.val_mask
-        else:
-            mask = graph.test_mask
-        edge_index = torch.stack(graph.edge_index)
-        edge_index, edge_types = edge_index[:, mask], graph.edge_attr[mask]
+    def loss(self, graph):
+        edge_index = graph.edge_index
+        edge_types = graph.edge_attr
 
         self.get_edge_set(edge_index, edge_types)
         batch_edges, batch_attr, samples, rels, labels = sampling_edge_uniform(
@@ -153,14 +147,5 @@ def predict(self, graph):
         indices = torch.arange(0, self.num_nodes).to(device)
         x = self.emb(indices)
         output = self.model(graph, x)
-        mrr, hits = cal_mrr(
-            output,
-            self.rel_weight.weight,
-            graph.edge_index,
-            graph.edge_attr,
-            scoring=self.scoring,
-            protocol="raw",
-            batch_size=500,
-            hits=[1, 3, 10],
-        )
-        return mrr, hits
+
+        return output, self.rel_weight.weight
diff --git a/cogdl/models/nn/self_auxiliary_task.py b/cogdl/models/nn/self_auxiliary_task.py
index b42ce762..cbb4472a 100644
--- a/cogdl/models/nn/self_auxiliary_task.py
+++ b/cogdl/models/nn/self_auxiliary_task.py
@@ -11,24 +11,18 @@
 from tqdm import tqdm
 
 from .. import register_model
-from cogdl.data import Dataset
-from cogdl.utils import dropout_adj
+from cogdl.utils.transform import dropout_adj
 from cogdl.models.nn.gcn import TKipfGCN
 from cogdl.models.self_supervised_model import SelfSupervisedGenerativeModel
 from cogdl.trainers.self_supervised_trainer import SelfSupervisedJointTrainer
 
 
 class SSLTask:
-    def __init__(self, graph, device):
-        self.graph = graph
-        self.edge_index = graph.edge_index_train if hasattr(graph, "edge_index_train") else graph.edge_index
-        self.num_nodes = graph.num_nodes
-        self.num_edges = graph.num_edges
-        self.features = graph.x
+    def __init__(self, device):
         self.device = device
         self.cached_edges = None
 
-    def transform_data(self):
+    def transform_data(self, graph):
         raise NotImplementedError
 
     def make_loss(self, embeddings):
@@ -36,44 +30,48 @@ def make_loss(self, embeddings):
 
 
 class EdgeMask(SSLTask):
-    def __init__(self, graph, hidden_size, mask_ratio, device):
-        super().__init__(graph, device)
+    def __init__(self, hidden_size, mask_ratio, device):
+        super().__init__(device)
         self.linear = nn.Linear(hidden_size, 2).to(device)
         self.mask_ratio = mask_ratio
 
-    def transform_data(self):
-        if self.cached_edges is None:
-            row, col = self.edge_index
-            edges = torch.stack([row, col])
-            perm = np.random.permutation(self.num_edges)
-            preserve_nnz = int(self.num_edges * (1 - self.mask_ratio))
-            masked = perm[preserve_nnz:]
-            preserved = perm[:preserve_nnz]
-            self.masked_edges = edges[:, masked]
-            self.cached_edges = edges[:, preserved]
-            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).to(self.device)
-            self.graph.edge_index = self.cached_edges
-
-        return self.graph.to(self.device)
+    def transform_data(self, graph):
+        device = graph.x.device
+        num_edges = graph.num_edges
+        # if self.cached_edges is None:
+        row, col = graph.edge_index
+        edges = torch.stack([row, col])
+        perm = np.random.permutation(num_edges)
+        preserve_nnz = int(num_edges * (1 - self.mask_ratio))
+        masked = perm[preserve_nnz:]
+        preserved = perm[:preserve_nnz]
+        self.masked_edges = edges[:, masked]
+        self.cached_edges = edges[:, preserved]
+        mask_num = len(masked)
+        self.neg_edges = self.neg_sample(mask_num, graph.edge_index)
+        self.pseudo_labels = torch.cat([torch.ones(mask_num), torch.zeros(mask_num)]).long().to(device)
+        self.node_pairs = torch.cat([self.masked_edges, self.neg_edges], 1).to(device)
+
+        graph.edge_index = self.cached_edges
+        return graph
 
     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 = torch.stack(self.edge_index).t().cpu().numpy()
+    def neg_sample(self, edge_num, graph):
+        edge_index = graph.edge_index
+        num_nodes = graph.num_nodes
+        edges = torch.stack(edge_index).t().cpu().numpy()
         exclude = set([(_[0], _[1]) for _ in list(edges)])
-        itr = self.sample(exclude)
+        itr = self.sample(exclude, num_nodes)
         sampled = [next(itr) for _ in range(edge_num)]
         return torch.tensor(sampled).t()
 
-    def sample(self, exclude):
+    def sample(self, exclude, num_nodes):
         while True:
-            t = tuple(np.random.randint(0, self.num_nodes, 2))
+            t = tuple(np.random.randint(0, num_nodes, 2))
             if t[0] != t[1] and t not in exclude:
                 exclude.add(t)
                 exclude.add((t[1], t[0]))
@@ -82,23 +80,25 @@ def sample(self, exclude):
 
 class AttributeMask(SSLTask):
     def __init__(self, graph, hidden_size, train_mask, mask_ratio, device):
-        super().__init__(graph, device)
+        super().__init__(device)
         self.linear = nn.Linear(hidden_size, graph.x.shape[1]).to(device)
-        self.unlabeled = np.array([i for i in range(self.num_nodes) if not train_mask[i]])
         self.cached_features = None
         self.mask_ratio = mask_ratio
 
-    def transform_data(self):
-        if self.cached_features is None:
-            perm = np.random.permutation(self.unlabeled)
-            mask_nnz = int(self.num_nodes * self.mask_ratio)
-            self.masked_nodes = perm[:mask_nnz]
-            self.cached_features = self.features.clone()
-            self.cached_features[self.masked_nodes] = torch.zeros(self.features.shape[1])
-            self.pseudo_labels = self.features[self.masked_nodes].to(self.device)
-            self.graph.features = self.cached_features
-
-        return self.graph.to(self.device)
+    def transform_data(self, graph):
+        # if self.cached_features is None:
+        device = graph.x.device
+        x_feat = graph.x
+
+        num_nodes = graph.num_nodes
+        unlabelled = torch.where(~graph.train_mask)[0]
+        perm = np.random.permutation(unlabelled)
+        mask_nnz = int(num_nodes * self.mask_ratio)
+        self.masked_nodes = perm[:mask_nnz]
+        x_feat[self.masked_nodes] = 0
+        self.pseudo_labels = x_feat[self.masked_nodes].to(device)
+        graph.x = x_feat
+        return graph
 
     def make_loss(self, embeddings):
         embeddings = self.linear(embeddings[self.masked_nodes])
@@ -107,8 +107,8 @@ def make_loss(self, embeddings):
 
 
 class PairwiseDistance(SSLTask):
-    def __init__(self, graph, hidden_size, class_split, sampling, dropedge_rate, num_centers, device):
-        super().__init__(graph, device)
+    def __init__(self, hidden_size, class_split, sampling, dropedge_rate, num_centers, device):
+        super().__init__(device)
         self.nclass = len(class_split) + 1
         self.class_split = class_split
         self.max_distance = self.class_split[self.nclass - 2][1]
@@ -116,15 +116,19 @@ def __init__(self, graph, hidden_size, class_split, sampling, dropedge_rate, num
         self.dropedge_rate = dropedge_rate
         self.num_centers = num_centers
         self.linear = nn.Linear(hidden_size, self.nclass).to(device)
-        self.get_distance()
+        self.get_distance_cache = False
+
+    def get_distance(self, graph):
+        num_nodes = graph.num_nodes
+        num_edges = graph.num_edges
+        edge_index = graph.edge_index
 
-    def get_distance(self):
         if self.sampling:
             self.dis_node_pairs = [[] for i in range(self.nclass)]
-            node_idx = random.sample(range(self.num_nodes), self.num_centers)
+            node_idx = random.sample(range(num_nodes), self.num_centers)
             adj = sp.coo_matrix(
-                (np.ones(self.num_edges), (self.edge_index[0].cpu().numpy(), self.edge_index[1].cpu().numpy())),
-                shape=(self.num_nodes, self.num_nodes),
+                (np.ones(num_edges), (edge_index[0].cpu().numpy(), edge_index[1].cpu().numpy())),
+                shape=(num_nodes, num_nodes),
             ).tocsr()
 
             num_samples = tqdm(range(self.num_centers))
@@ -132,7 +136,7 @@ def get_distance(self):
                 num_samples.set_description(f"Generating node pairs {i:03d}")
                 idx = node_idx[i]
                 queue = [idx]
-                dis = -np.ones(self.num_nodes)
+                dis = -np.ones(num_nodes)
                 dis[idx] = 0
                 head = 0
                 tail = 0
@@ -174,17 +178,17 @@ def get_distance(self):
                         (self.dis_node_pairs[cur_class], np.array([[idx] * len(sampled), sampled]).transpose()), axis=0
                     )
             if self.class_split[0][1] == 2:
-                self.dis_node_pairs[0] = torch.stack(self.edge_index).cpu().numpy().transpose()
+                self.dis_node_pairs[0] = torch.stack(edge_index).cpu().numpy().transpose()
             num_per_class = np.min(np.array([len(dis) for dis in self.dis_node_pairs]))
             for i in range(self.nclass):
                 sampled = np.random.choice(np.arange(len(self.dis_node_pairs[i])), num_per_class, replace=False)
                 self.dis_node_pairs[i] = self.dis_node_pairs[i][sampled]
         else:
             G = nx.Graph()
-            G.add_edges_from(torch.stack(self.edge_index).cpu().numpy().transpose())
+            G.add_edges_from(torch.stack(edge_index).cpu().numpy().transpose())
 
             path_length = dict(nx.all_pairs_shortest_path_length(G, cutoff=self.max_distance))
-            distance = -np.ones((self.num_nodes, self.num_nodes), dtype=np.int)
+            distance = -np.ones((num_nodes, num_nodes), dtype=np.int)
             for u, p in path_length.items():
                 for v, d in p.items():
                     distance[u][v] = d - 1
@@ -201,9 +205,12 @@ def get_distance(self):
             np.random.shuffle(tmp)
             self.dis_node_pairs.append(tmp)
 
-    def transform_data(self):
-        self.graph.edge_index, _ = dropout_adj(edge_index=self.edge_index, drop_rate=self.dropedge_rate)
-        return self.graph.to(self.device)
+    def transform_data(self, graph):
+        if not self.get_distance_cache:
+            self.get_distance(graph)
+            self.get_distance_cache = True
+        graph.edge_index, _ = dropout_adj(edge_index=graph.edge_index, drop_rate=self.dropedge_rate)
+        return graph
 
     def make_loss(self, embeddings, sample=True, k=4000):
         node_pairs, pseudo_labels = self.sample(sample, k)
@@ -231,18 +238,26 @@ def sample(self, sample, k):
 
 
 class Distance2Clusters(SSLTask):
-    def __init__(self, graph, hidden_size, num_clusters, device):
-        super().__init__(graph, device)
+    def __init__(self, hidden_size, num_clusters, device):
+        super().__init__(device)
         self.num_clusters = num_clusters
         self.linear = nn.Linear(hidden_size, num_clusters).to(device)
-        self.gen_cluster_info()
+        self.gen_cluster_info_cache = False
+
+    def transform_data(self, graph):
+        if not self.gen_cluster_info_cache:
+            self.gen_cluster_info_cache(graph)
+            self.gen_cluster_info_cache = True
+
+        return graph
 
-    def transform_data(self):
-        return self.graph.to(self.device)
+    def gen_cluster_info(self, graph, use_metis=False):
+        edge_index = graph.edge_index
+        num_nodes = graph.num_nodes
+        x = graph.x
 
-    def gen_cluster_info(self, use_metis=False):
         G = nx.Graph()
-        G.add_edges_from(torch.stack(self.edge_index).cpu().numpy().transpose())
+        G.add_edges_from(torch.stack(edge_index).cpu().numpy().transpose())
         if use_metis:
             import metis
 
@@ -250,14 +265,14 @@ def gen_cluster_info(self, use_metis=False):
         else:
             from sklearn.cluster import KMeans
 
-            clustering = KMeans(n_clusters=self.num_clusters, random_state=0).fit(self.features.cpu())
+            clustering = KMeans(n_clusters=self.num_clusters, random_state=0).fit(x.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))
+        self.distance_vec = np.zeros((num_nodes, self.num_clusters))
         for i in range(self.num_clusters):
             subgraph = G.subgraph(node_clusters[i])
             center = None
@@ -276,19 +291,22 @@ def make_loss(self, embeddings):
 
 
 class PairwiseAttrSim(SSLTask):
-    def __init__(self, graph, hidden_size, k, device):
-        super().__init__(graph, device)
+    def __init__(self, hidden_size, k, device):
+        super().__init__(device)
         self.k = k
         self.linear = nn.Linear(hidden_size, 1).to(self.device)
-        self.get_attr_sim()
+        self.get_attr_sim_cache = False
+
+    def get_avg_distance(self, graph, idx_sorted, k, sampled):
+        edge_index = graph.edge_index
+        num_nodes = graph.num_nodes
 
-    def get_avg_distance(self, idx_sorted, k, sampled):
         self.G = nx.Graph()
-        self.G.add_edges_from(torch.stack(self.edge_index).cpu().numpy().transpose())
+        self.G.add_edges_from(torch.stack(edge_index).cpu().numpy().transpose())
         avg_min = 0
         avg_max = 0
         avg_sampled = 0
-        for i in range(self.num_nodes):
+        for i in range(num_nodes):
             distance = dict(nx.shortest_path_length(self.G, source=i))
             sum = 0
             num = 0
@@ -297,7 +315,7 @@ def get_avg_distance(self, idx_sorted, k, sampled):
                     sum += distance[node]
                     num += 1
             if num:
-                avg_min += sum / num / self.num_nodes
+                avg_min += sum / num / num_nodes
             sum = 0
             num = 0
             for node in idx_sorted[i, -k - 1 :]:
@@ -305,7 +323,7 @@ def get_avg_distance(self, idx_sorted, k, sampled):
                     sum += distance[node]
                     num += 1
             if num:
-                avg_max += sum / num / self.num_nodes
+                avg_max += sum / num / num_nodes
             sum = 0
             num = 0
             for node in idx_sorted[i, sampled]:
@@ -313,18 +331,21 @@ def get_avg_distance(self, idx_sorted, k, sampled):
                     sum += distance[node]
                     num += 1
             if num:
-                avg_sampled += sum / num / self.num_nodes
+                avg_sampled += sum / num / num_nodes
         return avg_min, avg_max, avg_sampled
 
-    def get_attr_sim(self):
+    def get_attr_sim(self, graph):
+        x = graph.x
+        num_nodes = graph.num_nodes
+
         from sklearn.metrics.pairwise import cosine_similarity
 
-        sims = cosine_similarity(self.features.numpy())
+        sims = cosine_similarity(x.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):
+        sampled = self.sample(self.k, num_nodes)
+        for i in range(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]])
@@ -332,7 +353,7 @@ def get_attr_sim(self):
                 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))
+            % (self.get_avg_distance(graph, idx_sorted, self.k, sampled))
         )
         self.node_pairs = self.node_pairs.long().to(self.device)
         self.pseudo_labels = self.pseudo_labels.float().to(self.device)
@@ -340,11 +361,13 @@ def get_attr_sim(self):
     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)
+            sampled.append(int(random.random() * (num_nodes - self.k * 2)) + self.k)
         return np.array(sampled)
 
-    def transform_data(self):
-        return self.graph.to(self.device)
+    def transform_data(self, graph):
+        if not self.get_attr_sim_cache:
+            self.get_attr_sim(graph)
+        return graph
 
     def make_loss(self, embeddings):
         node_pairs = self.node_pairs
@@ -352,90 +375,91 @@ def make_loss(self, embeddings):
         return F.mse_loss(output, self.pseudo_labels, reduction="mean")
 
 
-@register_model("self_auxiliary_task")
-class SelfAuxiliaryTask(SelfSupervisedGenerativeModel):
-    @classmethod
-    def build_model_from_args(cls, args):
-        return cls(args)
-
-    @staticmethod
-    def add_args(parser):
-        """Add model-specific arguments to the parser."""
-        # fmt: off
-        parser.add_argument("--num-features", type=int)
-        parser.add_argument("--num-classes", type=int)
-        parser.add_argument("--auxiliary-task", type=str)
-        parser.add_argument("--num-layers", type=int, default=2)
-        parser.add_argument("--hidden-size", type=int, default=512)
-        parser.add_argument("--dropout", type=float, default=0.5)
-        parser.add_argument("--residual", action="store_true")
-        parser.add_argument("--norm", type=str, default=None)
-        parser.add_argument("--activation", type=str, default="relu")
-        parser.add_argument("--sampling", action="store_true")
-        parser.add_argument("--dropedge-rate", type=float, default=0.0)
-        parser.add_argument("--mask-ratio", type=float, default=0.1)
-        # fmt: on
-
-    def __init__(self, args):
-        super().__init__()
-        self.device = args.device_id[0] if not args.cpu else "cpu"
-        self.auxiliary_task = args.auxiliary_task
-        self.hidden_size = args.hidden_size
-        self.sampling = args.sampling
-        self.dropedge_rate = args.dropedge_rate
-        self.mask_ratio = args.mask_ratio
-        self.gcn = TKipfGCN(
-            args.num_features,
-            args.hidden_size,
-            args.num_classes,
-            args.num_layers,
-            args.dropout,
-            args.activation,
-            args.residual,
-            args.norm,
-        )
-        self.agent = None
-
-    def generate_virtual_labels(self, data):
-        if self.auxiliary_task == "edgemask":
-            self.agent = EdgeMask(data, self.hidden_size, self.mask_ratio, self.device)
-        elif self.auxiliary_task == "attributemask":
-            self.agent = AttributeMask(data, self.hidden_size, data.train_mask, self.mask_ratio, self.device)
-        elif self.auxiliary_task == "pairwise-distance":
-            self.agent = PairwiseDistance(
-                data,
-                self.hidden_size,
-                [(1, 2), (2, 3), (3, 5)],
-                self.sampling,
-                self.dropedge_rate,
-                256,
-                self.device,
-            )
-        elif self.auxiliary_task == "distance2clusters":
-            self.agent = Distance2Clusters(data, self.hidden_size, 30, self.device)
-        elif self.auxiliary_task == "pairwise-attr-sim":
-            self.agent = PairwiseAttrSim(data, self.hidden_size, 5, self.device)
-        else:
-            raise Exception(
-                "auxiliary task must be edgemask, pairwise-distance, distance2clusters, distance2clusters++ or pairwise-attr-sim"
-            )
-
-    def transform_data(self):
-        return self.agent.transform_data()
-
-    def self_supervised_loss(self, data):
-        embed = self.gcn.embed(data)
-        return self.agent.make_loss(embed)
-
-    def forward(self, data):
-        return self.gcn.forward(data)
-
-    def embed(self, data):
-        return self.gcn.embed(data)
-
-    def get_parameters(self):
-        return list(self.gcn.parameters()) + list(self.agent.linear.parameters())
-
-    @staticmethod
-    def get_trainer(args):
-        return SelfSupervisedJointTrainer
+#
+# @register_model("self_auxiliary_task")
+# class SelfAuxiliaryTask(SelfSupervisedGenerativeModel):
+#     @classmethod
+#     def build_model_from_args(cls, args):
+#         return cls(args)
+#
+#     @staticmethod
+#     def add_args(parser):
+#         """Add model-specific arguments to the parser."""
+#         # fmt: off
+#         parser.add_argument("--num-features", type=int)
+#         parser.add_argument("--num-classes", type=int)
+#         parser.add_argument("--auxiliary-task", type=str)
+#         parser.add_argument("--num-layers", type=int, default=2)
+#         parser.add_argument("--hidden-size", type=int, default=512)
+#         parser.add_argument("--dropout", type=float, default=0.5)
+#         parser.add_argument("--residual", action="store_true")
+#         parser.add_argument("--norm", type=str, default=None)
+#         parser.add_argument("--activation", type=str, default="relu")
+#         parser.add_argument("--sampling", action="store_true")
+#         parser.add_argument("--dropedge-rate", type=float, default=0.0)
+#         parser.add_argument("--mask-ratio", type=float, default=0.1)
+#         # fmt: on
+#
+#     def __init__(self, args):
+#         super().__init__()
+#         self.device = args.device_id[0] if not args.cpu else "cpu"
+#         self.auxiliary_task = args.auxiliary_task
+#         self.hidden_size = args.hidden_size
+#         self.sampling = args.sampling
+#         self.dropedge_rate = args.dropedge_rate
+#         self.mask_ratio = args.mask_ratio
+#         self.gcn = TKipfGCN(
+#             args.num_features,
+#             args.hidden_size,
+#             args.num_classes,
+#             args.num_layers,
+#             args.dropout,
+#             args.activation,
+#             args.residual,
+#             args.norm,
+#         )
+#         self.agent = None
+#
+#     def generate_virtual_labels(self, data):
+#         if self.auxiliary_task == "edgemask":
+#             self.agent = EdgeMask(data, self.hidden_size, self.mask_ratio, self.device)
+#         elif self.auxiliary_task == "attributemask":
+#             self.agent = AttributeMask(data, self.hidden_size, data.train_mask, self.mask_ratio, self.device)
+#         elif self.auxiliary_task == "pairwise-distance":
+#             self.agent = PairwiseDistance(
+#                 data,
+#                 self.hidden_size,
+#                 [(1, 2), (2, 3), (3, 5)],
+#                 self.sampling,
+#                 self.dropedge_rate,
+#                 256,
+#                 self.device,
+#             )
+#         elif self.auxiliary_task == "distance2clusters":
+#             self.agent = Distance2Clusters(data, self.hidden_size, 30, self.device)
+#         elif self.auxiliary_task == "pairwise-attr-sim":
+#             self.agent = PairwiseAttrSim(data, self.hidden_size, 5, self.device)
+#         else:
+#             raise Exception(
+#                 "auxiliary task must be edgemask, pairwise-distance, distance2clusters, distance2clusters++ or pairwise-attr-sim"
+#             )
+#
+#     def transform_data(self):
+#         return self.agent.transform_data()
+#
+#     def self_supervised_loss(self, data):
+#         embed = self.gcn.embed(data)
+#         return self.agent.make_loss(embed)
+#
+#     def forward(self, data):
+#         return self.gcn.forward(data)
+#
+#     def embed(self, data):
+#         return self.gcn.embed(data)
+#
+#     def get_parameters(self):
+#         return list(self.gcn.parameters()) + list(self.agent.linear.parameters())
+#
+#     @staticmethod
+#     def get_trainer(args):
+#         return SelfSupervisedJointTrainer
diff --git a/cogdl/operators/spmm.py b/cogdl/operators/spmm.py
index 948622c5..d2ff5b4b 100644
--- a/cogdl/operators/spmm.py
+++ b/cogdl/operators/spmm.py
@@ -7,6 +7,7 @@
 
 # SPMM
 
+
 try:
     spmm = load(
         name="spmm",
@@ -69,6 +70,7 @@ def backward(ctx, grad_out):
 except Exception:
     pass
 
+
 class ActSPMMFunction(torch.autograd.Function):
     @staticmethod
     def forward(ctx, rowptr, colind, feat, edge_weight_csr=None, sym=False):
diff --git a/cogdl/options.py b/cogdl/options.py
index 2ae4aca7..67d4c979 100644
--- a/cogdl/options.py
+++ b/cogdl/options.py
@@ -124,7 +124,7 @@ def parse_args_and_arch(parser, args):
             if hasattr(DATASET_REGISTRY[dataset], "add_args"):
                 DATASET_REGISTRY[dataset].add_args(parser)
 
-    if "trainer" in args and args.trainer is not None:
+    if "controller" in args and args.trainer is not None:
         if try_import_trainer(args.trainer):
             if hasattr(TRAINER_REGISTRY[args.trainer], "add_args"):
                 TRAINER_REGISTRY[args.trainer].add_args(parser)
diff --git a/cogdl/runner/__init__.py b/cogdl/runner/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/cogdl/runner/controller/__init__.py b/cogdl/runner/controller/__init__.py
new file mode 100644
index 00000000..d8848b17
--- /dev/null
+++ b/cogdl/runner/controller/__init__.py
@@ -0,0 +1,2 @@
+from .data_controller import DataController
+from .training_controller import TrainingController
diff --git a/cogdl/runner/controller/data_controller.py b/cogdl/runner/controller/data_controller.py
new file mode 100644
index 00000000..624d794d
--- /dev/null
+++ b/cogdl/runner/controller/data_controller.py
@@ -0,0 +1,49 @@
+import torch
+from cogdl.data import DataLoader
+
+from cogdl.wrappers.data_wrapper.base_data_wrapper import OnLoadingWrapper
+
+
+class DataController(object):
+    def __init__(self, world_size: int = 1, distributed: bool = False):
+        self.world_size = world_size
+        self.distributed = distributed
+
+    def distributed_dataloader(self, dataloader: DataLoader, dataset, rank):
+        # TODO: just a toy implementation
+        assert isinstance(dataloader, DataLoader)
+
+        args, kwargs = dataloader.get_parameters()
+        sampler = torch.utils.data.distributed.DistributedSampler(dataset, num_replicas=self.world_size, rank=rank)
+        kwargs["sampler"] = sampler
+        dataloader = dataloader.__class__(*args, **kwargs)
+        return dataloader
+
+    def prepare_data_wrapper(self, dataset_w, rank=0):
+        if self.distributed:
+            train_loader = dataset_w.train_wrapper()
+            assert isinstance(train_loader, DataLoader)
+            train_loader = self.distributed_dataloader(train_loader, dataset=dataset_w.get_train_dataset(), rank=rank)
+            train_wrapper = OnLoadingWrapper(train_loader, dataset_w.train_transform)
+            dataset_w.prepare_val_data()
+            dataset_w.prepare_test_data()
+            dataset_w.set_train_data(train_wrapper)
+            return dataset_w
+        else:
+            dataset_w.pre_transform()
+            dataset_w.prepare_training_data()
+            dataset_w.prepare_val_data()
+            dataset_w.prepare_test_data()
+            return dataset_w
+
+    def training_proc_per_stage(self, dataset_w, rank=0):
+        if dataset_w.__refresh_per_epoch__():
+            if self.distributed:
+                train_loader = dataset_w.train_wrapper()
+                assert isinstance(train_loader, DataLoader)
+                train_loader = self.distributed_dataloader(train_loader, dataset=dataset_w.get_dataset(), rank=rank)
+                train_wrapper = OnLoadingWrapper(train_loader, dataset_w.train_transform)
+                dataset_w.__train_data = train_wrapper
+            else:
+                dataset_w.prepare_training_data()
+        return dataset_w
diff --git a/cogdl/runner/controller/training_controller.py b/cogdl/runner/controller/training_controller.py
new file mode 100644
index 00000000..be334129
--- /dev/null
+++ b/cogdl/runner/controller/training_controller.py
@@ -0,0 +1,68 @@
+import os
+import logging
+
+import torch
+import torch.multiprocessing as mp
+
+from typing import List, Optional
+from cogdl.wrappers import ModelWrapper
+
+log = logging.getLogger(__name__)
+
+
+class TrainingController(object):
+    def __init__(self, device_ids: Optional[List[int]], dist: str = "ddp", backend: str = "nccl"):
+        self.device_ids = device_ids
+        self.backend = backend
+
+    def init_controller(self):
+        if self.backend == "ddp":
+            pass
+        elif self.backend == "dp":
+            pass
+        else:
+            raise NotImplementedError
+
+    def setup(self, model_w: ModelWrapper):
+        mp.spawn(self.new_process, args=(), nprocs=2)
+
+    def new_process(self) -> None:
+        pass
+
+    def init_ddp(self, global_rank: Optional[int], world_size: Optional[int]) -> None:
+        # TODO: this code is duplicated in DDP and DDPSpawn, make this a function
+        global_rank = global_rank if global_rank is not None else self.global_rank()
+        world_size = world_size if world_size is not None else self.world_size()
+        os.environ["MASTER_ADDR"] = self.master_address()
+        os.environ["MASTER_PORT"] = str(self.master_port())
+
+        if not torch.distributed.is_initialized():
+            log.info(f"initializing ddp: GLOBAL_RANK: {global_rank}, MEMBER: {global_rank + 1}/{world_size}")
+            torch.distributed.init_process_group(self.backend, rank=global_rank, world_size=world_size)
+
+    @property
+    def torch_distributed_backend(self):
+        torch_backend = os.getenv("PL_TORCH_DISTRIBUTED_BACKEND")
+        if torch_backend is None:
+            torch_backend = "nccl" if self.on_gpu else "gloo"
+        return torch_backend
+
+    def local_rank(self) -> int:
+        return int(os.environ["LOCAL_RANK"])
+
+    def master_address(self) -> str:
+        if "MASTER_ADDR" not in os.environ:
+            # rank_zero_warn("MASTER_ADDR environment variable is not defined. Set as localhost")
+            os.environ["MASTER_ADDR"] = "127.0.0.1"
+        log.debug(f"MASTER_ADDR: {os.environ['MASTER_ADDR']}")
+        master_address = os.environ.get("MASTER_ADDR")
+        return master_address
+
+    def master_port(self) -> int:
+        if "MASTER_PORT" not in os.environ:
+            # rank_zero_warn("MASTER_PORT environment variable is not defined. Set as 12910")
+            os.environ["MASTER_PORT"] = "12910"
+        log.debug(f"MASTER_PORT: {os.environ['MASTER_PORT']}")
+
+        port = int(os.environ.get("MASTER_PORT"))
+        return port
diff --git a/cogdl/runner/embed_runner.py b/cogdl/runner/embed_runner.py
new file mode 100644
index 00000000..96a31f51
--- /dev/null
+++ b/cogdl/runner/embed_runner.py
@@ -0,0 +1,58 @@
+import os
+import numpy as np
+import time
+from typing import Optional
+import torch
+
+
+class EmbeddingTrainer(object):
+    def __init__(
+        self,
+        save_embedding_path: Optional[str] = None,
+        load_embedding_path: Optional[str] = None,
+    ):
+        self.save_embedding_path = save_embedding_path
+        self.default_embedding_dir = "./embeddings"
+        self.load_embedding_path = load_embedding_path
+
+    def run(self, model_w, dataset_w):
+        self.prepare_data_wrapper(dataset_w)
+        if self.load_embedding_path is not None:
+            embedding = np.load(self.load_embedding_path)
+            return self.test(model_w, dataset_w, embedding)
+
+        if self.save_embedding_path is None:
+            cur_time = time.strftime("%m-%d_%H.%M.%S", time.localtime())
+            name = f"{model_w.wrapped_model.__class__.__name__}_{cur_time}.emb"
+            self.save_embedding_path = os.path.join(self.default_embedding_dir, name)
+            os.makedirs(self.default_embedding_dir, exist_ok=True)
+        embeddings = self.train(model_w, dataset_w)
+        self.save_embedding(embeddings)
+        return self.test(model_w, dataset_w, embeddings)
+
+    def prepare_data_wrapper(self, dataset_w):
+        dataset_w.pre_transform()
+        dataset_w.prepare_training_data()
+        dataset_w.prepare_val_data()
+        dataset_w.prepare_test_data()
+
+    def train(self, model_w, dataset_w):
+        dataset_w.pre_transform()
+        train_data = dataset_w.on_train_wrapper()
+        embeddings = []
+        for batch in train_data:
+            embeddings.append(model_w.train_step(batch))
+        # embeddings = model_w.train_step(train_data)
+        assert len(embeddings) == 1
+        embeddings = embeddings[0]
+        return embeddings
+
+    def test(self, model_w, dataset_w, embeddings):
+        labels = next(dataset_w.on_test_wrapper())
+        if torch.is_tensor(labels):
+            labels = labels.cpu().numpy()
+        result = model_w.test_step((embeddings, labels))
+        return result
+
+    def save_embedding(self, embeddings):
+        np.save(self.save_embedding_path, embeddings)
diff --git a/cogdl/runner/experiments.py b/cogdl/runner/experiments.py
new file mode 100644
index 00000000..538da7f3
--- /dev/null
+++ b/cogdl/runner/experiments.py
@@ -0,0 +1,262 @@
+import copy
+import itertools
+import os
+import inspect
+from collections import defaultdict, namedtuple
+
+import torch
+import optuna
+from tabulate import tabulate
+
+from cogdl.utils import set_random_seed, tabulate_results, init_operator_configs
+from cogdl.configs import BEST_CONFIGS
+from cogdl.models import build_model
+from cogdl.datasets import build_dataset
+from cogdl.wrappers import fetch_model_wrapper, fetch_data_wrapper
+from cogdl.runner.options import get_default_args
+from cogdl.runner.runner import Trainer
+
+
+class AutoML(object):
+    """
+    Args:
+        func_search: function to obtain hyper-parameters to search
+    """
+
+    def __init__(self, args):
+        self.func_search = args.func_search
+        self.metric = args.metric if hasattr(args, "metric") else None
+        self.n_trials = args.n_trials if hasattr(args, "n_trials") else 3
+        self.best_value = None
+        self.best_params = None
+        self.default_params = args
+
+    def _objective(self, trials):
+        params = copy.deepcopy(self.default_params)
+        cur_params = self.func_search(trials)
+        print(cur_params)
+        for key, value in cur_params.items():
+            params.__setattr__(key, value)
+        result_dict = raw_experiment(args=params)
+        result_list = list(result_dict.values())[0]
+        item = result_list[0]
+        key = self.metric
+        if key is None:
+            for _key in item.keys():
+                if "Val" in _key or "val" in _key:
+                    key = _key
+                    break
+            if key is None:
+                raise KeyError("Unable to find validation metrics")
+        val = [result[key] for result in result_list]
+        mean = sum(val) / len(val)
+
+        if self.best_value is None or mean > self.best_value:
+            self.best_value = mean
+            self.best_params = cur_params
+            self.best_results = result_list
+
+        return mean
+
+    def run(self):
+        study = optuna.create_study(direction="maximize")
+        study.optimize(self._objective, n_trials=self.n_trials, n_jobs=1)
+        print(study.best_params)
+        return self.best_results
+
+
+def examine_link_prediction(args, dataset):
+    if "link_prediction" in args.mw:
+        args.num_entities = dataset.data.num_nodes
+        # args.num_entities = len(torch.unique(self.data.edge_index))
+        if dataset.data.edge_attr is not None:
+            args.num_rels = len(torch.unique(dataset.data.edge_attr))
+            args.monitor = "mrr"
+        else:
+            args.monitor = "auc"
+    return args
+
+
+def set_best_config(args):
+    configs = BEST_CONFIGS[args.task]
+    if args.model not in configs:
+        return args
+    configs = configs[args.model]
+    for key, value in configs["general"].items():
+        args.__setattr__(key, value)
+    if args.dataset not in configs:
+        return args
+    for key, value in configs[args.dataset].items():
+        args.__setattr__(key, value)
+    return args
+
+
+def train(args):
+    print(
+        f""" 
+    |---------------------------------------------------{'-' * (len(args.mw) + len(args.dw))}|
+     *** Using `{args.mw}` ModelWrapper and `{args.dw}` DataWrapper 
+    |---------------------------------------------------{'-' * (len(args.mw) + len(args.dw))}|"""
+    )
+    set_random_seed(args.seed)
+
+    if getattr(args, "use_best_config", False):
+        args = set_best_config(args)
+
+    # setup dataset and specify `num_features` and `num_classes` for model
+    args.monitor = "val_acc"
+    dataset = build_dataset(args)
+    args.num_features = dataset.num_features
+    args.num_classes = dataset.num_classes
+    if hasattr(dataset, "num_nodes"):
+        args.num_nodes = dataset.num_nodes
+    if hasattr(dataset, "num_edge"):
+        args.num_edge = dataset.num_edge
+    if hasattr(args, "unsup") and args.unsup:
+        args.num_classes = args.hidden_size
+    else:
+        args.num_classes = dataset.num_classes
+
+    mw_class = fetch_model_wrapper(args.mw)
+    dw_class = fetch_data_wrapper(args.dw)
+
+    if mw_class is None:
+        raise NotImplementedError("`model wrapper(--mw)` must be specified.")
+
+    if dw_class is None:
+        raise NotImplementedError("`data wrapper(--dw)` must be specified.")
+
+    data_wrapper_args = dict()
+    model_wrapper_args = dict()
+    # unworthy code: share `args` between model and dataset_wrapper
+    for key in inspect.signature(dw_class).parameters.keys():
+        if hasattr(args, key) and key != "dataset":
+            data_wrapper_args[key] = getattr(args, key)
+    # unworthy code: share `args` between model and model_wrapper
+    for key in inspect.signature(mw_class).parameters.keys():
+        if hasattr(args, key) and key != "model":
+            model_wrapper_args[key] = getattr(args, key)
+
+    args = examine_link_prediction(args, dataset)
+
+    # setup model
+    model = build_model(args)
+    # specify configs for optimizer
+    optimizer_cfg = dict(
+        lr=args.lr,
+        weight_decay=args.weight_decay,
+        n_warmup_steps=args.n_warmup_steps,
+        max_epoch=args.max_epoch,
+        batch_size=args.batch_size if hasattr(args, "batch_size") else 0,
+    )
+
+    if hasattr(args, "hidden_size"):
+        optimizer_cfg["hidden_size"] = args.hidden_size
+
+    # setup model_wrapper
+    if "embedding" in args.mw:
+        model_wrapper = mw_class(model, **model_wrapper_args)
+    else:
+        model_wrapper = mw_class(model, optimizer_cfg, **model_wrapper_args)
+    # setup data_wrapper
+    dataset_wrapper = dw_class(dataset, **data_wrapper_args)
+
+    save_embedding_path = args.emb_path if hasattr(args, "emb_path") else None
+    os.makedirs("./checkpoints", exist_ok=True)
+    os.makedirs("./checkpoints", exist_ok=True)
+
+    # logger = fetch_logger(args.logger)
+
+    # setup controller
+    trainer = Trainer(
+        max_epoch=args.max_epoch,
+        device_ids=args.devices,
+        cpu=args.cpu,
+        save_embedding_path=save_embedding_path,
+        cpu_inference=args.cpu_inference,
+        # monitor=args.monitor,
+        progress_bar=args.progress_bar,
+        distributed_training=args.distributed,
+        checkpoint_path=args.checkpoint_path,
+        patience=args.patience,
+        logger=args.logger,
+        log_path=args.log_path,
+        project=args.project,
+    )
+
+    # Go!!!
+    result = trainer.run(model_wrapper, dataset_wrapper)
+
+    return result
+
+
+def gen_variants(**items):
+    Variant = namedtuple("Variant", items.keys())
+    return itertools.starmap(Variant, itertools.product(*items.values()))
+
+
+def variant_args_generator(args, variants):
+    """Form variants as group with size of num_workers"""
+    for variant in variants:
+        args.dataset, args.model, args.seed = variant
+        yield copy.deepcopy(args)
+
+
+def output_results(results_dict, tablefmt="github"):
+    variant = list(results_dict.keys())[0]
+    col_names = ["Variant"] + list(results_dict[variant][-1].keys())
+    tab_data = tabulate_results(results_dict)
+
+    print(tabulate(tab_data, headers=col_names, tablefmt=tablefmt))
+
+
+def raw_experiment(args):
+    init_operator_configs(args)
+
+    variants = list(gen_variants(dataset=args.dataset, model=args.model, seed=args.seed))
+
+    results_dict = defaultdict(list)
+    results = [train(args) for args in variant_args_generator(args, variants)]
+    for variant, result in zip(variants, results):
+        results_dict[variant[:-1]].append(result)
+
+    tablefmt = args.tablefmt if hasattr(args, "tablefmt") else "github"
+    output_results(results_dict, tablefmt)
+
+    return results_dict
+
+
+def auto_experiment(args):
+    variants = list(gen_variants(dataset=args.dataset, model=args.model))
+
+    results_dict = defaultdict(list)
+    for variant in variants:
+        args.model = [variant.model]
+        args.dataset = [variant.dataset]
+        tool = AutoML(args)
+        results_dict[variant[:]] = tool.run()
+
+    tablefmt = args.tablefmt if hasattr(args, "tablefmt") else "github"
+    print("\nFinal results:\n")
+    output_results(results_dict, tablefmt)
+
+    return results_dict
+
+
+def experiment(dataset, model, **kwargs):
+    if isinstance(dataset, str):
+        dataset = [dataset]
+    if isinstance(model, str):
+        model = [model]
+    if "args" not in kwargs:
+        args = get_default_args(dataset=dataset, model=model, **kwargs)
+    else:
+        args = kwargs["args"]
+        for key, value in kwargs.items():
+            if key != "args":
+                args.__setattr__(key, value)
+
+    if "func_search" in kwargs:
+        return auto_experiment(args)
+
+    return raw_experiment(args)
diff --git a/cogdl/runner/options.py b/cogdl/runner/options.py
new file mode 100644
index 00000000..19d44654
--- /dev/null
+++ b/cogdl/runner/options.py
@@ -0,0 +1,180 @@
+import sys
+import argparse
+import copy
+import warnings
+
+from cogdl.datasets import try_import_dataset
+from cogdl.models import MODEL_REGISTRY, try_import_model
+from cogdl.wrappers import fetch_data_wrapper, fetch_model_wrapper
+from cogdl.utils import build_args_from_dict
+from cogdl.wrappers.default_match import get_wrappers_name
+
+
+def add_args(args: list):
+    parser = argparse.ArgumentParser()
+    if "lr" in args:
+        parser.add_argument("--lr", default=0.01, type=float)
+    if "max_epoch" in args:
+        parser.add_argument("--max-epoch", default=500, type=int)
+
+
+def add_arguments(args: list):
+    parser = argparse.ArgumentParser()
+    for item in args:
+        name, _type, default = item
+        parser.add_argument(f"--{name}", default=default, type=_type)
+    return parser
+
+
+def get_parser():
+    parser = argparse.ArgumentParser(conflict_handler="resolve")
+    # fmt: off
+    parser.add_argument("--seed", default=[1], type=int, nargs="+", metavar="N",
+                        help="pseudo random number generator seed")
+    parser.add_argument("--max-epoch", default=500, type=int)
+    parser.add_argument("--patience", type=int, default=100)
+    parser.add_argument("--lr", default=0.01, type=float)
+    parser.add_argument("--weight-decay", default=5e-4, type=float)
+    parser.add_argument("--n-warmup-steps", type=int, default=0)
+
+    parser.add_argument("--checkpoint-path", type=str, default="./checkpoints/model.pt", help="path to save model")
+    parser.add_argument("--logger", type=str, default=None)
+    parser.add_argument("--log-path", type=str, default=".", help="path to save logs")
+    parser.add_argument("--project", type=str, default="cogdl-exp", help="project name for wandb")
+
+    parser.add_argument("--use-best-config", action="store_true", help="use best config")
+    parser.add_argument("--unsup", action="store_true")
+
+    parser.add_argument("--devices", default=[0], type=int, nargs="+", help="which GPU to use")
+    parser.add_argument("--cpu", action="store_true", help="use CPU instead of CUDA")
+    parser.add_argument("--cpu-inference", action="store_true", help="do validation and test in cpu")
+    # parser.add_argument("--monitor", type=str, default="val_acc")
+    parser.add_argument("--distributed", action="store_true")
+    parser.add_argument("--progress-bar", type=str, default="epoch")
+    parser.add_argument("--local_rank", type=int, default=0)
+    parser.add_argument("--master-port", type=int, default=13425)
+    parser.add_argument("--master-addr", type=str, default="localhost")
+    # fmt: on
+    return parser
+
+
+def add_data_wrapper_args(parser):
+    group = parser.add_argument_group("Data wrapper configuration")
+    # fmt: off
+    group.add_argument("--dw", "-t", type=str, default=None, metavar="DWRAPPER", required=False,
+                       help="Data Wrapper")
+    # fmt: on
+    return group
+
+
+def add_model_wrapper_args(parser):
+    group = parser.add_argument_group("Trainer configuration")
+    # fmt: off
+    group.add_argument("--mw", type=str, default=None, metavar="MWRAPPER", required=False,
+                       help="Model Wrapper")
+    # fmt: on
+    return group
+
+
+def add_dataset_args(parser):
+    group = parser.add_argument_group("Dataset and data loading")
+    # fmt: off
+    group.add_argument("--dataset", "-dt", metavar="DATASET", nargs="+", required=True,
+                       help="Dataset")
+    # fmt: on
+    return group
+
+
+def add_model_args(parser):
+    group = parser.add_argument_group("Model configuration")
+    # fmt: off
+    group.add_argument("--model", "-m", metavar="MODEL", nargs="+", required=True,
+                       help="Model Architecture")
+    # fmt: on
+    return group
+
+
+def get_training_parser():
+    parser = get_parser()
+    add_dataset_args(parser)
+    add_model_args(parser)
+    add_data_wrapper_args(parser)
+    add_model_wrapper_args(parser)
+    return parser
+
+
+def get_display_data_parser():
+    parser = get_parser()
+    add_dataset_args(parser)
+    parser.add_argument("--depth", default=3, type=int)
+
+    return parser
+
+
+def get_download_data_parser():
+    parser = get_parser()
+    add_dataset_args(parser)
+
+    return parser
+
+
+def get_default_args(dataset, model, **kwargs):
+    if not isinstance(dataset, list):
+        dataset = [dataset]
+    if not isinstance(model, list):
+        model = [model]
+    sys.argv = [sys.argv[0], "-m"] + model + ["-dt"] + dataset
+    if "mw" in kwargs:
+        sys.argv += ["--mw"] + [kwargs["mw"]]
+    if "dw" in kwargs:
+        sys.argv += ["--dw"] + [kwargs["dw"]]
+
+    # The parser doesn"t know about specific args, so we parse twice.
+    parser = get_training_parser()
+    args, _ = parser.parse_known_args()
+    args = parse_args_and_arch(parser, args)
+    for key, value in kwargs.items():
+        args.__setattr__(key, value)
+    return args
+
+
+def get_diff_args(args1, args2):
+    d1 = copy.deepcopy(args1.__dict__)
+    d2 = args2.__dict__
+    for k in d2.keys():
+        d1.pop(k, None)
+    return build_args_from_dict(d1)
+
+
+def parse_args_and_arch(parser, args):
+    # Add *-specific args to parser.
+    for model in args.model:
+        if try_import_model(model):
+            MODEL_REGISTRY[model].add_args(parser)
+
+    for dataset in args.dataset:
+        try_import_dataset(dataset)
+
+    if len(args.model) > 1:
+        warnings.warn("Please ensure that models could use the same model wrapper!")
+    default_wrappers = get_wrappers_name(args.model[0])
+    if default_wrappers is not None:
+        mw, dw = default_wrappers
+    else:
+        mw, dw = None, None
+
+    if args.dw is not None:
+        dw = args.dw
+    if hasattr(fetch_data_wrapper(dw), "add_args"):
+        fetch_data_wrapper(dw).add_args(parser)
+
+    if args.mw is not None:
+        mw = args.mw
+    if hasattr(fetch_model_wrapper(mw), "add_args"):
+        fetch_model_wrapper(mw).add_args(parser)
+
+    # Parse a second time.
+    args = parser.parse_args()
+    args.mw = mw
+    args.dw = dw
+    return args
diff --git a/cogdl/runner/runner.py b/cogdl/runner/runner.py
new file mode 100644
index 00000000..2c6afd6d
--- /dev/null
+++ b/cogdl/runner/runner.py
@@ -0,0 +1,436 @@
+import copy
+from typing import Optional
+import numpy as np
+from tqdm import tqdm
+import os
+
+import torch
+import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel
+import torch.multiprocessing as mp
+
+from cogdl.wrappers.data_wrapper.base_data_wrapper import DataWrapper
+from cogdl.wrappers.model_wrapper.base_model_wrapper import ModelWrapper, EmbeddingModelWrapper
+from cogdl.runner.runner_utils import evaluation_comp, load_model, save_model, ddp_end, ddp_after_epoch, Printer
+from cogdl.runner.embed_runner import EmbeddingTrainer
+from cogdl.runner.controller import DataController
+from cogdl.loggers import build_logger
+from cogdl.data import Graph
+
+from cogdl.datasets import build_dataset_from_name
+
+
+def move_to_device(batch, device):
+    if isinstance(batch, list) or isinstance(batch, tuple):
+        if isinstance(batch, tuple):
+            batch = list(batch)
+        for i, x in enumerate(batch):
+            if torch.is_tensor(x):
+                batch[i] = x.to(device)
+            elif isinstance(x, Graph):
+                x.to(device)
+    elif torch.is_tensor(batch) or isinstance(batch, Graph):
+        batch = batch.to(device)
+    elif hasattr(batch, "apply_to_device"):
+        batch.apply_to_device(device)
+    return batch
+
+
+def clip_grad_norm(params, max_norm):
+    """Clips gradient norm."""
+    if max_norm > 0:
+        return torch.nn.utils.clip_grad_norm_(params, max_norm)
+    else:
+        return torch.sqrt(sum(p.grad.data.norm() ** 2 for p in params if p.grad is not None))
+
+
+class Trainer(object):
+    def __init__(
+        self,
+        max_epoch: int,
+        nstage: int = 1,
+        cpu: bool = False,
+        checkpoint_path: str = "./checkpoints/checkpoint.pt",
+        device_ids: Optional[list] = None,
+        distributed_training: bool = False,
+        distributed_inference: bool = False,
+        master_addr: str = "localhost",
+        master_port: int = 10086,
+        # monitor: str = "val_acc",
+        early_stopping: bool = True,
+        patience: int = 100,
+        eval_step: int = 1,
+        save_embedding_path: Optional[str] = None,
+        cpu_inference: bool = False,
+        progress_bar: str = "epoch",
+        clip_grad_norm: float = 5.0,
+        logger: str = None,
+        log_path: str = "./runs",
+        project: str = "cogdl-exp",
+    ):
+        self.max_epoch = max_epoch
+        self.nstage = nstage
+        self.patience = patience
+        self.early_stopping = early_stopping
+        self.eval_step = eval_step
+        self.monitor = None
+        self.evaluation_metric = None
+        self.progress_bar = progress_bar
+
+        self.cpu = cpu
+        self.devices, self.world_size = self.set_device(device_ids)
+        self.checkpoint_path = checkpoint_path
+
+        self.distributed_training = distributed_training
+        self.distributed_inference = distributed_inference
+        # if self.world_size <= 1:
+        #     self.distributed_training = False
+        #     self.distributed_inference = False
+        self.master_addr = master_addr
+        self.master_port = master_port
+
+        self.cpu_inference = cpu_inference
+
+        self.on_train_batch_transform = None
+        self.on_eval_batch_transform = None
+        self.clip_grad_norm = clip_grad_norm
+
+        self.save_embedding_path = save_embedding_path
+
+        self.data_controller = DataController(world_size=self.world_size, distributed=self.distributed_training)
+
+        self.logger = build_logger(logger, log_path, project)
+
+        self.after_epoch_hooks = []
+        self.pre_epoch_hooks = []
+        self.training_end_hooks = []
+
+        if distributed_training:
+            self.register_training_end_hook(ddp_end)
+            self.register_out_epoch_hook(ddp_after_epoch)
+
+        self.eval_data_back_to_cpu = False
+
+    def register_in_epoch_hook(self, hook):
+        self.pre_epoch_hooks.append(hook)
+
+    def register_out_epoch_hook(self, hook):
+        self.after_epoch_hooks.append(hook)
+
+    def register_training_end_hook(self, hook):
+        self.training_end_hooks.append(hook)
+
+    def set_device(self, device_ids: Optional[list]):
+        """
+        Return: devices, world_size
+        """
+        if device_ids is None or self.cpu:
+            return [torch.device("cpu")], 0
+
+        if isinstance(device_ids, int) and device_ids > 0:
+            device_ids = [device_ids]
+        elif isinstance(device_ids, list):
+            pass
+        else:
+            raise ValueError("`device_id` has to be list of integers")
+        if len(device_ids) == 0:
+            return torch.device("cpu"), 0
+        else:
+            return [i for i in device_ids], len(device_ids)
+
+    def run(self, model_w: ModelWrapper, dataset_w: DataWrapper):
+        # for network/graph embedding models
+        if isinstance(model_w, EmbeddingModelWrapper):
+            return EmbeddingTrainer(self.save_embedding_path).run(model_w, dataset_w)
+
+        # for deep learning models
+        # set default loss_fn and evaluator for model_wrapper
+        # mainly for in-cogdl setting
+
+        model_w.default_loss_fn = dataset_w.get_default_loss_fn()
+        model_w.default_evaluator = dataset_w.get_default_evaluator()
+        model_w.set_evaluation_metric()
+
+        if self.distributed_training:  # and self.world_size > 1:
+            torch.multiprocessing.set_sharing_strategy("file_system")
+            self.dist_train(model_w, dataset_w)
+        else:
+            self.train(self.devices[0], model_w, dataset_w)
+        best_model_w = load_model(model_w, self.checkpoint_path).to(self.devices[0])
+
+        # disable `distributed` to inference once only
+        self.distributed_training = False
+        dataset_w.prepare_test_data()
+        final_val = self.validate(model_w, dataset_w, self.devices[0])
+        final_test = self.test(best_model_w, dataset_w, self.devices[0])
+
+        if "val_metric" in final_val:
+            final_val[f"val_{self.evaluation_metric}"] = final_val["val_metric"]
+            final_val.pop("val_metric")
+            final_val.pop("val_loss")
+        if "test_metric" in final_test:
+            final_test[f"test_{self.evaluation_metric}"] = final_test["test_metric"]
+            final_test.pop("test_metric")
+            final_test.pop("test_loss")
+        self.logger.note(final_test)
+        final_test.update(final_val)
+        print(final_test)
+        return final_test
+
+    def dist_train(self, model_w: ModelWrapper, dataset_w: DataWrapper):
+        mp.set_start_method("spawn", force=True)
+
+        device_count = torch.cuda.device_count()
+        if device_count < self.world_size:
+            size = device_count
+            print(f"Available device count ({device_count}) is less than world size ({self.world_size})")
+        else:
+            size = self.world_size
+
+        print(f"Let's using {size} GPUs.")
+
+        processes = []
+        for rank in range(size):
+            p = mp.Process(target=self.train, args=(rank, model_w, dataset_w))
+
+            p.start()
+            print(f"Process [{rank}] starts!")
+            processes.append(p)
+
+        for p in processes:
+            p.join()
+
+    def build_optimizer(self, model_w):
+        opt_wrap = model_w.setup_optimizer()
+        if isinstance(opt_wrap, list) or isinstance(opt_wrap, tuple):
+            assert len(opt_wrap) == 2
+            optimizers, lr_schedulars = opt_wrap
+        else:
+            optimizers = opt_wrap
+            lr_schedulars = None
+
+        if not isinstance(optimizers, list):
+            optimizers = [optimizers]
+        if lr_schedulars and not isinstance(lr_schedulars, list):
+            lr_schedulars = [lr_schedulars]
+        return optimizers, lr_schedulars
+
+    def initialize(self, model_w, rank=0, master_addr: str = "localhost", master_port: int = 10008):
+        if self.distributed_training:
+            os.environ["MASTER_ADDR"] = master_addr
+            os.environ["MASTER_PORT"] = str(master_port)
+            dist.init_process_group("nccl", rank=rank, world_size=self.world_size)
+            model_w = copy.deepcopy(model_w).to(rank)
+            model_w = DistributedDataParallel(model_w, device_ids=[rank])
+
+            module = model_w.module
+            model_w, model_ddp = module, model_w
+            return model_w, model_ddp
+        else:
+            return model_w.to(rank), None
+
+    def train(self, rank, model_w, dataset_w):
+        model_w, _ = self.initialize(model_w, rank=rank, master_addr=self.master_addr, master_port=self.master_port)
+        self.data_controller.prepare_data_wrapper(dataset_w, rank)
+        self.eval_data_back_to_cpu = dataset_w.data_back_to_cpu
+
+        optimizers, lr_schedulars = self.build_optimizer(model_w)
+
+        est = model_w.set_early_stopping()
+        if isinstance(est, str):
+            est_monitor = est
+            best_index, compare_fn = evaluation_comp(est_monitor)
+        else:
+            assert len(est) == 2
+            est_monitor, est_compare = est
+            best_index, compare_fn = evaluation_comp(est_monitor, est_compare)
+        self.monitor = est_monitor
+        self.evaluation_metric = model_w.evaluation_metric
+
+        # best_index, compare_fn = evaluation_comp(self.monitor)
+        best_model_w = None
+
+        patience = 0
+        best_epoch = 0
+        for stage in range(self.nstage):
+            with torch.no_grad():
+                pre_stage_out = model_w.pre_stage(stage, dataset_w)
+                dataset_w.pre_stage(stage, pre_stage_out)
+                self.data_controller.training_proc_per_stage(dataset_w, rank)
+
+            if self.progress_bar == "epoch":
+                epoch_iter = tqdm(range(self.max_epoch))
+                epoch_printer = Printer(epoch_iter.set_description, rank=rank, world_size=self.world_size)
+            else:
+                epoch_iter = range(self.max_epoch)
+                epoch_printer = Printer(print, rank=rank, world_size=self.world_size)
+
+            self.logger.start()
+            for epoch in epoch_iter:
+                print_str_dict = dict()
+                for hook in self.pre_epoch_hooks:
+                    hook(self)
+
+                # inductive setting ..
+                dataset_w.train()
+                train_loader = dataset_w.on_train_wrapper()
+                training_loss = self.training_step(model_w, train_loader, optimizers, lr_schedulars, rank)
+
+                print_str_dict["Epoch"] = epoch
+                print_str_dict["train_loss"] = training_loss
+
+                val_loader = dataset_w.on_val_wrapper()
+                if val_loader is not None and (epoch % self.eval_step) == 0:
+                    # inductive setting ..
+                    dataset_w.eval()
+                    # do validation in inference device
+                    val_result = self.validate(model_w, dataset_w, rank)
+                    # print(val_result)
+                    if val_result is not None:
+                        monitoring = val_result[self.monitor]
+                        if compare_fn(monitoring, best_index):
+                            best_index = monitoring
+                            best_epoch = epoch
+                            patience = 0
+                            best_model_w = copy.deepcopy(model_w)
+                        else:
+                            patience += 1
+                            if self.early_stopping and patience >= self.patience:
+                                break
+                        print_str_dict[f"val_{self.evaluation_metric}"] = monitoring
+
+                if self.distributed_training:
+                    if rank == 0:
+                        epoch_printer(print_str_dict)
+                        self.logger.note(print_str_dict, epoch)
+                else:
+                    epoch_printer(print_str_dict)
+                    self.logger.note(print_str_dict, epoch)
+
+                for hook in self.after_epoch_hooks:
+                    hook(self)
+
+            with torch.no_grad():
+                model_w.eval()
+                post_stage_out = model_w.post_stage(stage, dataset_w)
+                dataset_w.post_stage(stage, post_stage_out)
+
+            if best_model_w is None:
+                best_model_w = copy.deepcopy(model_w)
+
+        if self.distributed_training:
+            if rank == 0:
+                save_model(best_model_w.to("cpu"), self.checkpoint_path, best_epoch)
+                dist.barrier()
+            else:
+                dist.barrier()
+        else:
+            save_model(best_model_w.to("cpu"), self.checkpoint_path, best_epoch)
+
+        for hook in self.training_end_hooks:
+            hook(self)
+
+    def validate(self, model_w: ModelWrapper, dataset_w: DataWrapper, device):
+        # ------- distributed training ---------
+        if self.distributed_training:
+            return self.distributed_test(model_w, dataset_w.on_val_wrapper(), device, self.val_step)
+        # ------- distributed training ---------
+
+        model_w.eval()
+        if self.cpu_inference:
+            model_w.to("cpu")
+            _device = device
+        else:
+            _device = device
+
+        val_loader = dataset_w.on_val_wrapper()
+        result = self.val_step(model_w, val_loader, _device)
+
+        model_w.to(device)
+        return result
+
+    def test(self, model_w: ModelWrapper, dataset_w: DataWrapper, device):
+        # ------- distributed training ---------
+        if self.distributed_training:
+            return self.distributed_test(model_w, dataset_w.on_test_wrapper(), device, self.test_step)
+        # ------- distributed training ---------
+
+        model_w.eval()
+        if self.cpu_inference:
+            model_w.to("cpu")
+            _device = device
+        else:
+            _device = device
+
+        test_loader = dataset_w.on_test_wrapper()
+        result = self.test_step(model_w, test_loader, _device)
+
+        model_w.to(device)
+        return result
+
+    def distributed_test(self, model_w: ModelWrapper, loader, rank, fn):
+        model_w.eval()
+        # if rank == 0:
+        if dist.get_rank() == 0:
+            if self.cpu_inference:
+                model_w.to("cpu")
+                _device = "cpu"
+            else:
+                _device = rank
+            result = fn(model_w, loader, _device)
+            model_w.to(rank)
+
+            object_list = [result]
+        else:
+            object_list = [None]
+        dist.broadcast_object_list(object_list, src=0)
+        return object_list[0]
+
+    def training_step(self, model_w, train_loader, optimizers, lr_schedulars, device):
+        model_w.train()
+        losses = []
+
+        if self.progress_bar == "iteration":
+            train_loader = tqdm(train_loader)
+
+        for batch in train_loader:
+            # batch = batch.to(device)
+            batch = move_to_device(batch, device)
+            loss = model_w.on_train_step(batch)
+
+            for optimizer in optimizers:
+                optimizer.zero_grad()
+            loss.backward()
+
+            torch.nn.utils.clip_grad_norm_(model_w.parameters(), self.clip_grad_norm)
+
+            for optimizer in optimizers:
+                optimizer.step()
+
+            losses.append(loss.item())
+        if lr_schedulars is not None:
+            for lr_schedular in lr_schedulars:
+                lr_schedular.step()
+        return np.mean(losses)
+
+    @torch.no_grad()
+    def val_step(self, model_w, val_loader, device):
+        model_w.eval()
+        for batch in val_loader:
+            # batch = batch.to(device)
+            batch = move_to_device(batch, device)
+            model_w.on_val_step(batch)
+            if self.eval_data_back_to_cpu:
+                move_to_device(batch, "cpu")
+        return model_w.collect_notes()
+
+    @torch.no_grad()
+    def test_step(self, model_w, test_loader, device):
+        model_w.eval()
+        for batch in test_loader:
+            # batch = batch.to(device)
+            batch = move_to_device(batch, device)
+            model_w.on_test_step(batch)
+            if self.eval_data_back_to_cpu:
+                move_to_device(batch, "cpu")
+        return model_w.collect_notes()
diff --git a/cogdl/runner/runner_utils.py b/cogdl/runner/runner_utils.py
new file mode 100644
index 00000000..320982b5
--- /dev/null
+++ b/cogdl/runner/runner_utils.py
@@ -0,0 +1,89 @@
+from typing import Dict
+
+import numpy as np
+
+import torch
+import torch.distributed as dist
+
+
+def merge_batch_indexes(outputs: list):
+    assert len(outputs) > 0
+    keys = list(outputs[0].keys())
+
+    results = dict()
+    for key in keys:
+        values = [x[key] for x in outputs]
+        if key.endswith("loss"):
+            results[key] = sum(values).item() / len(values)
+        elif key.endswith("eval_index"):
+            if len(values) > 1:
+                val = torch.cat(values, dim=0)
+                val = val.sum(0)
+            else:
+                val = values[0]
+            fp = val[0]
+            all_ = val.sum()
+
+            prefix = key[: key.find("eval_index")]
+            if val.shape[0] == 2:
+                _key = prefix + "acc"
+            else:
+                _key = prefix + "f1"
+            results[_key] = (fp / all_).item()
+        else:
+            results[key] = sum(values)
+    return results
+
+
+def bigger_than(x, y):
+    return x >= y
+
+
+def smaller_than(x, y):
+    return x <= y
+
+
+def evaluation_comp(monitor, compare="<"):
+    if "loss" in monitor or compare == "<":
+        return np.inf, smaller_than
+    else:
+        return 0, bigger_than
+
+
+def save_model(model, path, epoch):
+    print(f"Saving {epoch}-th model to {path} ...")
+    torch.save(model.state_dict(), path)
+
+
+def load_model(model, path):
+    print(f"Loading model from {path} ...")
+    model.load_state_dict(torch.load(path))
+    return model
+
+
+def ddp_after_epoch(*args):
+    dist.barrier()
+
+
+def ddp_end(*args):
+    dist.barrier()
+    dist.destroy_process_group()
+
+
+class Printer(object):
+    def __init__(self, print_fn, rank=0, world_size=1):
+        self.printer = print_fn
+        self.to_print = (world_size <= 1) or rank == 0 or rank == "cpu"
+
+    def __call__(self, k_v: Dict):
+        if self.to_print:
+            assert "Epoch" in k_v
+            out = f"Epoch: {k_v['Epoch']}"
+            k_v.pop("Epoch")
+
+            for k, v in k_v.items():
+                if isinstance(v, float):
+                    out += f", {k}: {v: .4f}"
+                else:
+                    out += f", {k}: {v}"
+            self.printer(out)
diff --git a/cogdl/tasks/base_task.py b/cogdl/tasks/base_task.py
index 2b4de0a2..ee49b1eb 100644
--- a/cogdl/tasks/base_task.py
+++ b/cogdl/tasks/base_task.py
@@ -64,7 +64,7 @@ def set_evaluator(self, dataset):
         self.evaluator = dataset.get_evaluator()
 
     def get_trainer(self, args):
-        if hasattr(args, "trainer") and args.trainer is not None:
+        if hasattr(args, "controller") and args.trainer is not None:
             if "self_auxiliary_task" in args.trainer and not hasattr(self.model, "embed"):
                 raise ValueError("Model ({}) must implement embed method".format(args.model))
             return build_trainer(args)
diff --git a/cogdl/tasks/link_prediction.py b/cogdl/tasks/link_prediction.py
index 333885de..6f9ac7f1 100644
--- a/cogdl/tasks/link_prediction.py
+++ b/cogdl/tasks/link_prediction.py
@@ -18,6 +18,7 @@
 from tqdm import tqdm
 
 from . import BaseTask, register_task
+from ..data import Graph
 
 
 def save_model(model, optimizer, save_variable_list, args):
@@ -203,10 +204,11 @@ def __init__(self, args, dataset=None, model=None):
         self.model.set_device(self.device)
 
     def train(self):
-        G = nx.Graph()
-        G.add_edges_from(self.train_data)
+        train_data = torch.from_numpy(np.array(self.train_data).transpose())
+        G = Graph(edge_index=train_data)
         embeddings = self.model.train(G)
 
+        G = G.to_networkx()
         embs = dict()
         for vid, node in enumerate(G.nodes()):
             embs[node] = embeddings[vid]
diff --git a/cogdl/tasks/node_classification.py b/cogdl/tasks/node_classification.py
index df7d6296..34763dc7 100644
--- a/cogdl/tasks/node_classification.py
+++ b/cogdl/tasks/node_classification.py
@@ -21,6 +21,7 @@ def add_args(parser: argparse.ArgumentParser):
         # fmt: off
         parser.add_argument("--missing-rate", type=int, default=0, help="missing rate, from 0 to 100")
         parser.add_argument("--inference", action="store_true")
+        parser.add_argument("--actnn", action="store_true")
         # fmt: on
 
     def __init__(
@@ -100,9 +101,9 @@ def train(self):
             for epoch in epoch_iter:
                 self._train_step()
                 acc, losses = self._test_step()
-                train_acc = acc["train"]
-                val_acc = acc["val"]
-                val_loss = losses["val"]
+                train_acc = acc["train"].item()
+                val_acc = acc["val"].item()
+                val_loss = losses["val"].item()
                 epoch_iter.set_description(
                     f"Epoch: {epoch:03d}, Train: {train_acc:.4f}, Val: {val_acc:.4f}, ValLoss:{val_loss: .4f}"
                 )
@@ -113,7 +114,7 @@ def train(self):
                         # best_loss = val_loss
                         best_score = val_acc
                         best_model = copy.deepcopy(self.model)
-                    min_loss = np.min((min_loss, val_loss.cpu()))
+                    min_loss = np.min((min_loss, val_loss))
                     max_score = np.max((max_score, val_acc))
                     patience = 0
                 else:
@@ -124,7 +125,7 @@ def train(self):
             print(f"Valid accurracy = {best_score: .4f}")
             self.model = best_model
         acc, _ = self._test_step(post=True)
-        val_acc, test_acc = acc["val"], acc["test"]
+        val_acc, test_acc = acc["val"].item(), acc["test"].item()
 
         print(f"Test accuracy = {test_acc:.4f}")
 
diff --git a/cogdl/tasks/unsupervised_node_classification.py b/cogdl/tasks/unsupervised_node_classification.py
index 2478b73c..54a47ada 100644
--- a/cogdl/tasks/unsupervised_node_classification.py
+++ b/cogdl/tasks/unsupervised_node_classification.py
@@ -64,7 +64,7 @@ def __init__(self, args, dataset=None, model=None):
         self.enhance = args.enhance
         self.training_percents = args.training_percents
         self.args = args
-        self.is_weighted = self.data.edge_attr is not None
+        self.is_weighted = self.data.edge_weight is not None
         self.device = "cpu" if not torch.cuda.is_available() or args.cpu else args.device_id[0]
 
         self.trainer = self.get_trainer(args)
@@ -110,12 +110,12 @@ def train(self):
                 if self.is_weighted:
                     edges, weight = (
                         edge_index.t().tolist(),
-                        self.data.edge_attr.tolist(),
+                        self.data.edge_weight.tolist(),
                     )
-                    G.add_weighted_edges_from([(edges[i][0], edges[i][1], weight[0][i]) for i in range(len(edges))])
+                    G.add_weighted_edges_from([(edges[i][0], edges[i][1], weight[i]) for i in range(len(edges))])
                 else:
                     G.add_edges_from(edge_index.t().tolist())
-                embeddings = self.model.train(G)
+                embeddings = self.model(G)
                 if self.enhance is not None:
                     embeddings = self.enhance_emb(G, embeddings)
                 # Map node2id
diff --git a/cogdl/trainers/__init__.py b/cogdl/trainers/__init__.py
index b57419a4..252cc272 100644
--- a/cogdl/trainers/__init__.py
+++ b/cogdl/trainers/__init__.py
@@ -8,7 +8,7 @@
 
 def register_trainer(name):
     """
-    New universal trainer types can be added to cogdl with the :func:`register_trainer`
+    New universal controller types can be added to cogdl with the :func:`register_trainer`
     function decorator.
 
     For example::
@@ -23,7 +23,7 @@ class SelfAuxiliaryTaskTrainer(BaseModel):
 
     def register_trainer_cls(cls):
         if name in TRAINER_REGISTRY:
-            raise ValueError("Cannot register duplicate universal trainer ({})".format(name))
+            raise ValueError("Cannot register duplicate universal controller ({})".format(name))
         if not issubclass(cls, BaseTrainer):
             raise ValueError("Model ({}: {}) must extend BaseTrainer".format(name, cls.__name__))
         TRAINER_REGISTRY[name] = cls
@@ -38,7 +38,7 @@ def try_import_trainer(trainer):
         if trainer in SUPPORTED_TRAINERS:
             importlib.import_module(SUPPORTED_TRAINERS[trainer])
         else:
-            print(f"Failed to import {trainer} trainer.")
+            print(f"Failed to import {trainer} controller.")
             return False
     return True
 
diff --git a/cogdl/trainers/base_trainer.py b/cogdl/trainers/base_trainer.py
index fe4e3060..178eaa7b 100644
--- a/cogdl/trainers/base_trainer.py
+++ b/cogdl/trainers/base_trainer.py
@@ -20,7 +20,7 @@ def __init__(self, args=None):
     @classmethod
     @abstractmethod
     def build_trainer_from_args(cls, args):
-        """Build a new trainer instance."""
+        """Build a new controller instance."""
         raise NotImplementedError("Trainers must implement the build_trainer_from_args method")
 
     @abstractmethod
diff --git a/cogdl/trainers/daegc_trainer.py b/cogdl/trainers/daegc_trainer.py
index 9191c029..d75391c2 100644
--- a/cogdl/trainers/daegc_trainer.py
+++ b/cogdl/trainers/daegc_trainer.py
@@ -10,6 +10,7 @@
 
 class DAEGCTrainer(BaseTrainer):
     def __init__(self, args):
+        super(DAEGCTrainer, self).__init__()
         self.num_clusters = args.num_clusters
         self.max_epoch = args.max_epoch
         self.lr = args.lr
diff --git a/cogdl/trainers/distributed_sampled_trainer.py b/cogdl/trainers/distributed_sampled_trainer.py
index 268df34b..97d34e18 100644
--- a/cogdl/trainers/distributed_sampled_trainer.py
+++ b/cogdl/trainers/distributed_sampled_trainer.py
@@ -174,6 +174,7 @@ def __init__(self, args):
         super(DistributedClusterGCNTrainer, self).__init__(args)
 
     def fit(self, *args, **kwargs):
+        print("MyFit")
         return super(DistributedClusterGCNTrainer, self).dist_fit(*args, **kwargs)
 
     def build_dataloader(self, dataset, rank):
@@ -193,7 +194,7 @@ def build_dataloader(self, dataset, rank):
         )
 
         if torch.__version__.split("+")[0] < "1.7.1":
-            settings.pop("persistent_wo rkers")
+            settings.pop("persistent_workers")
 
         data.train()
         train_loader = ClusteredLoader(
diff --git a/cogdl/trainers/distributed_trainer.py b/cogdl/trainers/distributed_trainer.py
index a332dec1..76337b8b 100644
--- a/cogdl/trainers/distributed_trainer.py
+++ b/cogdl/trainers/distributed_trainer.py
@@ -164,7 +164,7 @@ def train(model, dataset, args, rank, evaluator, loss_fn):
 class DistributedClusterGCNTrainer(BaseTrainer):
     @staticmethod
     def add_args(parser: argparse.ArgumentParser):
-        """Add trainer-specific arguments to the parser."""
+        """Add controller-specific arguments to the parser."""
         # fmt: off
         parser.add_argument("--n-cluster", type=int, default=1000)
         parser.add_argument("--batch-size", type=int, default=20)
diff --git a/cogdl/trainers/sampled_trainer.py b/cogdl/trainers/sampled_trainer.py
index 4912d26a..a434f4cc 100644
--- a/cogdl/trainers/sampled_trainer.py
+++ b/cogdl/trainers/sampled_trainer.py
@@ -93,7 +93,7 @@ def train(self):
 class SAINTTrainer(SampledTrainer):
     @staticmethod
     def add_args(parser: argparse.ArgumentParser):
-        """Add trainer-specific arguments to the parser."""
+        """Add controller-specific arguments to the parser."""
         # fmt: off
         SampledTrainer.add_args(parser)
         parser.add_argument("--eval-cpu", action="store_true")
@@ -197,7 +197,7 @@ class NeighborSamplingTrainer(SampledTrainer):
 
     @staticmethod
     def add_args(parser: argparse.ArgumentParser):
-        """Add trainer-specific arguments to the parser."""
+        """Add controller-specific arguments to the parser."""
         # fmt: off
         SampledTrainer.add_args(parser)
         # fmt: on
@@ -288,7 +288,7 @@ def build_trainer_from_args(cls, args):
 class ClusterGCNTrainer(SampledTrainer):
     @staticmethod
     def add_args(parser: argparse.ArgumentParser):
-        """Add trainer-specific arguments to the parser."""
+        """Add controller-specific arguments to the parser."""
         # fmt: off
         SampledTrainer.add_args(parser)
         parser.add_argument("--n-cluster", type=int, default=1000)
diff --git a/cogdl/trainers/self_supervised_trainer.py b/cogdl/trainers/self_supervised_trainer.py
index f662476e..15d4926b 100644
--- a/cogdl/trainers/self_supervised_trainer.py
+++ b/cogdl/trainers/self_supervised_trainer.py
@@ -16,7 +16,7 @@
 class SelfSupervisedBaseTrainer(BaseTrainer):
     @staticmethod
     def add_args(parser: argparse.ArgumentParser):
-        """Add trainer-specific arguments to the parser."""
+        """Add controller-specific arguments to the parser."""
         # fmt: off
         parser.add_argument('--subgraph-sampling', action='store_true')
         parser.add_argument('--sample-size', type=int, default=8192)
@@ -46,7 +46,7 @@ def fit(self, model, dataset):
 class SelfSupervisedJointTrainer(SelfSupervisedBaseTrainer):
     @staticmethod
     def add_args(parser: argparse.ArgumentParser):
-        """Add trainer-specific arguments to the parser."""
+        """Add controller-specific arguments to the parser."""
         # fmt: off
         SelfSupervisedBaseTrainer.add_args(parser)
         parser.add_argument('--alpha', default=10, type=float)
@@ -131,7 +131,7 @@ def _test_step(self, split="train"):
 class SelfSupervisedPretrainer(SelfSupervisedBaseTrainer):
     @staticmethod
     def add_args(parser: argparse.ArgumentParser):
-        """Add trainer-specific arguments to the parser."""
+        """Add controller-specific arguments to the parser."""
         # fmt: off
         SelfSupervisedBaseTrainer.add_args(parser)
         parser.add_argument('--alpha', default=1, type=float)
diff --git a/cogdl/utils/__init__.py b/cogdl/utils/__init__.py
index bb9c90c3..d604da2c 100644
--- a/cogdl/utils/__init__.py
+++ b/cogdl/utils/__init__.py
@@ -3,3 +3,4 @@
 from .sampling import *
 from .graph_utils import *
 from .spmm_utils import *
+from .transform import *
diff --git a/cogdl/utils/evaluator.py b/cogdl/utils/evaluator.py
index 6591b9a5..c1c420b7 100644
--- a/cogdl/utils/evaluator.py
+++ b/cogdl/utils/evaluator.py
@@ -1,7 +1,132 @@
+from typing import Union, Callable
+import numpy as np
+
 import torch
+import torch.nn as nn
+
 from sklearn.metrics import f1_score
 
 
+def setup_evaluator(metric: Union[str, Callable]):
+    if isinstance(metric, str):
+        metric = metric.lower()
+        if metric == "acc" or metric == "accuracy":
+            return Accuracy()
+        elif metric == "multilabel_microf1" or "microf1" or "micro_f1":
+            return MultiLabelMicroF1()
+        elif metric == "multiclass_microf1":
+            return MultiClassMicroF1()
+        else:
+            raise NotImplementedError
+    else:
+        return BaseEvaluator(metric)
+
+
+class BaseEvaluator(object):
+    def __init__(self, eval_func):
+        self.y_pred = list()
+        self.y_true = list()
+        self.eval_func = eval_func
+
+    def __call__(self, y_pred, y_true):
+        metric = self.eval_func(y_pred, y_true)
+        self.y_pred.append(y_pred.cpu())
+        self.y_true.append(y_true.cpu())
+        return metric
+
+    def clear(self):
+        self.y_pred = list()
+        self.y_true = list()
+
+    def evaluate(self):
+        if len(self.y_pred) > 0:
+            y_pred = torch.cat(self.y_pred, dim=0)
+            y_true = torch.cat(self.y_true, dim=0)
+            self.clean()
+            return self.eval_func(y_pred, y_true)
+        return 0
+
+
+class Accuracy(object):
+    def __init__(self, mini_batch=False):
+        super(Accuracy, self).__init__()
+        self.mini_batch = mini_batch
+        self.tp = list()
+        self.total = list()
+
+    def __call__(self, y_pred, y_true):
+        pred = (y_pred.argmax(1) == y_true).int()
+        tp = pred.sum().int()
+        total = pred.shape[0]
+        if torch.is_tensor(tp):
+            tp = tp.item()
+
+        # if self.mini_batch:
+        self.tp.append(tp)
+        self.total.append(total)
+
+        return tp / total
+
+    def evaluate(self):
+        if len(self.tp) > 0:
+            tp = np.sum(self.tp)
+            total = np.sum(self.total)
+            self.tp = list()
+            self.total = list()
+            return tp / total
+        return 0
+
+    def clear(self):
+        self.tp = list()
+        self.total = list()
+
+
+class MultiLabelMicroF1(Accuracy):
+    def __init__(self, mini_batch=False):
+        super(MultiLabelMicroF1, self).__init__(mini_batch)
+
+    def __call__(self, y_pred, y_true, sigmoid=False):
+        if sigmoid:
+            border = 0.5
+        else:
+            border = 0
+        y_pred[y_pred >= border] = 1
+        y_pred[y_pred < border] = 0
+        tp = (y_pred * y_true).sum().to(torch.float32).item()
+        fp = ((1 - y_true) * y_pred).sum().to(torch.float32).item()
+        fn = (y_true * (1 - y_pred)).sum().to(torch.float32).item()
+        total = tp + fp + fn
+
+        # if self.mini_batch:
+        self.tp.append(int(tp))
+        self.total.append(int(total))
+
+        if total == 0:
+            return 0
+        return float(tp / total)
+
+
+class MultiClassMicroF1(Accuracy):
+    def __init__(self, mini_batch=False):
+        super(MultiClassMicroF1, self).__init__(mini_batch)
+
+
+class CrossEntropyLoss(nn.Module):
+    def __call__(self, y_pred, y_true):
+        y_true = y_true.long()
+        y_pred = torch.nn.functional.log_softmax(y_pred, dim=-1)
+        return torch.nn.functional.nll_loss(y_pred, y_true)
+
+
+class BCEWithLogitsLoss(nn.Module):
+    def __call__(self, y_pred, y_true, reduction="mean"):
+        y_true = y_true.float()
+        loss = torch.nn.BCEWithLogitsLoss(reduction=reduction)(y_pred, y_true)
+        if reduction == "none":
+            loss = torch.sum(torch.mean(loss, dim=0))
+        return loss
+
+
 def multilabel_f1(y_pred, y_true, sigmoid=False):
     if sigmoid:
         y_pred[y_pred > 0.5] = 1
diff --git a/cogdl/utils/graph_utils.py b/cogdl/utils/graph_utils.py
index e70c498f..1a4ae0fd 100644
--- a/cogdl/utils/graph_utils.py
+++ b/cogdl/utils/graph_utils.py
@@ -191,50 +191,18 @@ def remove_self_loops(indices, values=None):
     return (row, col), values
 
 
-def filter_adj(row, col, edge_attr, mask):
-    return (row[mask], col[mask]), None if edge_attr is None else edge_attr[mask]
-
-
-def dropout_adj(
-    edge_index: Tuple,
-    edge_weight: Optional[torch.Tensor] = None,
-    drop_rate: float = 0.5,
-    renorm: Optional[str] = "sym",
-    training: bool = False,
-):
-    if not training or drop_rate == 0:
-        if edge_weight is None:
-            edge_weight = torch.ones(edge_index[0].shape[0], device=edge_index[0].device)
-        return edge_index, edge_weight
-
-    if drop_rate < 0.0 or drop_rate > 1.0:
-        raise ValueError("Dropout probability has to be between 0 and 1, " "but got {}".format(drop_rate))
-
-    row, col = edge_index
-    num_nodes = int(max(row.max(), col.max())) + 1
-    self_loop = row == col
-    mask = torch.full((row.shape[0],), 1 - drop_rate, dtype=torch.float, device=row.device)
-    mask = torch.bernoulli(mask).to(torch.bool)
-    mask = self_loop | mask
-    edge_index, edge_weight = filter_adj(row, col, edge_weight, mask)
-    if renorm == "sym":
-        edge_weight = symmetric_normalization(num_nodes, edge_index[0], edge_index[1])
-    elif renorm == "row":
-        edge_weight = row_normalization(num_nodes, edge_index[0], edge_index[1])
-    return edge_index, edge_weight
-
-
 def coalesce(row, col, value=None):
+    device = row.device
     if torch.is_tensor(row):
-        row = row.numpy()
+        row = row.cpu().numpy()
     if torch.is_tensor(col):
-        col = col.numpy()
+        col = col.cpu().numpy()
     indices = np.lexsort((col, row))
-    row = torch.from_numpy(row[indices]).long()
-    col = torch.from_numpy(col[indices]).long()
+    row = torch.from_numpy(row[indices]).long().to(device)
+    col = torch.from_numpy(col[indices]).long().to(device)
 
     num = col.shape[0] + 1
-    idx = torch.full((num,), -1, dtype=torch.long)
+    idx = torch.full((num,), -1, dtype=torch.long).to(device)
     max_num = max(row.max(), col.max()) + 100
     idx[1:] = (row + 1) * max_num + col
     mask = idx[1:] > idx[:-1]
@@ -243,7 +211,7 @@ def coalesce(row, col, value=None):
         return row, col, value
     row = row[mask]
     if value is not None:
-        _value = torch.zeros(row.shape[0], dtype=torch.float).to(row.device)
+        _value = torch.zeros(row.shape[0], dtype=torch.float).to(device)
         value = _value.scatter_add_(dim=0, src=value, index=col)
     col = col[mask]
     return row, col, value
diff --git a/cogdl/utils/index.py b/cogdl/utils/index.py
new file mode 100644
index 00000000..1ddb13ce
--- /dev/null
+++ b/cogdl/utils/index.py
@@ -0,0 +1,41 @@
+import torch
+from .spmm_utils import spmm
+
+
+@torch.no_grad()
+def homo_index(g, x):
+    with g.local_graph():
+        g.remove_self_loops()
+        neighbors = spmm(g, x)
+        deg = g.degrees()
+    isolated_nodes = deg == 0
+    diff = (x - neighbors).norm(2, dim=-1)
+    diff = diff.mean(1)
+    diff = diff[~isolated_nodes]
+    return torch.mean(diff)
+
+
+@torch.no_grad()
+def mad_index(g, x):
+    row, col = g.edge_index
+    self_loop = row == col
+    mask = ~self_loop
+    row = row[mask]
+    col = col[mask]
+
+    src, tgt = x[col], x[row]
+    sim = (src * tgt).sum(dim=1)
+    src_size = src.norm(p=2, dim=1)
+    tgt_size = tgt.norm(p=2, dim=1)
+    distance = 1 - sim / (src_size * tgt_size)
+
+    N = g.num_nodes
+
+    deg = g.degrees() - 1
+    out = torch.zeros((N,), dtype=torch.float, device=x.device)
+    out = out.scatter_add_(index=row, dim=0, src=distance)
+    deg_inv = deg.pow(-1)
+    deg_inv[torch.isinf(deg_inv)] = 1
+    dis = out * deg_inv
+    dis = dis[dis > 0]
+    return torch.mean(dis).item()
diff --git a/cogdl/utils/link_prediction_utils.py b/cogdl/utils/link_prediction_utils.py
index 64d9e343..5f0b372a 100644
--- a/cogdl/utils/link_prediction_utils.py
+++ b/cogdl/utils/link_prediction_utils.py
@@ -96,23 +96,13 @@ def predict(self, sub_emb, obj_emb, rel_emb):
 
 
 class GNNLinkPredict(nn.Module):
-    def __init__(self, score_func, dim):
+    def __init__(self):
         super(GNNLinkPredict, self).__init__()
         self.edge_set = None
-        self.score_func = score_func
-        if score_func == "distmult":
-            self.scoring = DistMultLayer()
-        elif score_func == "conve":
-            self.scoring = ConvELayer(dim)
-        else:
-            raise NotImplementedError
 
     def forward(self, graph):
         raise NotImplementedError
 
-    def get_score(self, heads, tails, rels):
-        return self.scoring(heads, tails, rels)
-
     def get_edge_set(self, edge_index, edge_types):
         if self.edge_set is None:
             edge_list = torch.stack((edge_index[0], edge_index[1], edge_types))
diff --git a/cogdl/utils/optimizer.py b/cogdl/utils/optimizer.py
new file mode 100644
index 00000000..37e9f5bf
--- /dev/null
+++ b/cogdl/utils/optimizer.py
@@ -0,0 +1,78 @@
+"""A wrapper class for optimizer """
+import numpy as np
+import torch.nn as nn
+
+
+class NoamOptimizer(nn.Module):
+    """A simple wrapper class for learning rate scheduling"""
+
+    def __init__(self, optimizer, d_model, n_warmup_steps, init_lr=None):
+        super(NoamOptimizer, self).__init__()
+        self._optimizer = optimizer
+        self.param_groups = optimizer.param_groups
+        self.n_warmup_steps = n_warmup_steps
+        self.n_current_steps = 0
+        self.init_lr = np.power(d_model, -0.5) if init_lr is None else init_lr / np.power(self.n_warmup_steps, -0.5)
+
+    def step(self):
+        """Step with the inner optimizer"""
+        self._update_learning_rate()
+        self._optimizer.step()
+
+    def zero_grad(self):
+        """Zero out the gradients by the inner optimizer"""
+        self._optimizer.zero_grad()
+
+    def _get_lr_scale(self):
+        return np.min(
+            [np.power(self.n_current_steps, -0.5), np.power(self.n_warmup_steps, -1.5) * self.n_current_steps]
+        )
+
+    def _update_learning_rate(self):
+        """ Learning rate scheduling per step """
+
+        self.n_current_steps += 1
+        lr = self.init_lr * self._get_lr_scale()
+
+        for param_group in self._optimizer.param_groups:
+            param_group["lr"] = lr
+
+
+class LinearOptimizer(nn.Module):
+    """A simple wrapper class for learning rate scheduling"""
+
+    def __init__(self, optimizer, n_warmup_steps, n_training_steps, init_lr=0.001):
+        super(LinearOptimizer, self).__init__()
+        self._optimizer = optimizer
+        self.param_groups = optimizer.param_groups
+        self.n_warmup_steps = n_warmup_steps
+        self.n_current_steps = 0
+        self.n_training_steps = n_training_steps
+        self.init_lr = init_lr
+
+    def step(self):
+        """Step with the inner optimizer"""
+        self._update_learning_rate()
+        self._optimizer.step()
+
+    def zero_grad(self):
+        """Zero out the gradients by the inner optimizer"""
+        self._optimizer.zero_grad()
+
+    def _get_lr_scale(self):
+        if self.n_current_steps < self.n_warmup_steps:
+            return float(self.n_current_steps) / float(max(1, self.n_warmup_steps))
+        return max(
+            0.0,
+            float(self.n_training_steps - self.n_current_steps)
+            / float(max(1, self.n_training_steps - self.n_warmup_steps)),
+        )
+
+    def _update_learning_rate(self):
+        """ Learning rate scheduling per step """
+
+        self.n_current_steps += 1
+        lr = self.init_lr * self._get_lr_scale()
+
+        for param_group in self._optimizer.param_groups:
+            param_group["lr"] = lr
diff --git a/cogdl/utils/rwalk/Makefile b/cogdl/utils/rwalk/Makefile
new file mode 100644
index 00000000..a29a43b6
--- /dev/null
+++ b/cogdl/utils/rwalk/Makefile
@@ -0,0 +1,14 @@
+CC?=gcc # Set compiler if CC is not set
+CFLAGS= -fopenmp -fPIC -O3 -D NDEBUG -Wall -Werror
+
+all: librwalk.so
+
+librwalk.so: rwalk.o
+	$(CC) $(CFLAGS) -shared -Wl,-soname,librwalk.so -o librwalk.so rwalk.o
+	rm rwalk.o
+
+rwalk.o: rwalk.c
+	$(CC) -c $(CFLAGS) rwalk.c -o rwalk.o
+
+clean :
+	rm -rf librwalk.so rwalk.o __pycache__
diff --git a/cogdl/utils/rwalk/__init__.py b/cogdl/utils/rwalk/__init__.py
new file mode 100644
index 00000000..8fbe5804
--- /dev/null
+++ b/cogdl/utils/rwalk/__init__.py
@@ -0,0 +1,55 @@
+"""
+
+"""
+import numpy as np
+import numpy.ctypeslib as npct
+from ctypes import c_float, c_int
+from os.path import dirname
+
+array_1d_int = npct.ndpointer(dtype=np.int32, ndim=1, flags="CONTIGUOUS")
+
+librwalk = npct.load_library("librwalk", dirname(__file__))
+
+# print("rwalk: Loading library from: {}".format(dirname(__file__)))
+# librwalk.random_walk.restype = None
+# librwalk.random_walk.argtypes = [array_1d_int, array_1d_int, c_int, c_int, c_int, c_int, c_int, array_1d_int]
+
+librwalk.random_walk.restype = None
+librwalk.random_walk.argtypes = [
+    array_1d_int,
+    array_1d_int,
+    array_1d_int,
+    c_int,
+    c_int,
+    c_int,
+    c_int,
+    c_int,
+    c_float,
+    array_1d_int,
+]
+
+
+def random_walk(nodes, ptr, neighs, num_walks=1, num_steps=1, nthread=-1, seed=111413, restarp_prop=0.0):
+    assert ptr.flags["C_CONTIGUOUS"]
+    assert neighs.flags["C_CONTIGUOUS"]
+    assert ptr.dtype == np.int32
+    assert neighs.dtype == np.int32
+    assert nodes.dtype == np.int32
+    n = nodes.size
+    walks = -np.ones((n * num_walks, (num_steps + 1)), dtype=np.int32, order="C")
+    assert walks.flags["C_CONTIGUOUS"]
+
+    librwalk.random_walk(
+        nodes,
+        ptr,
+        neighs,
+        n,
+        num_walks,
+        num_steps,
+        seed,
+        nthread,
+        restarp_prop,
+        np.reshape(walks, (walks.size,), order="C"),
+    )
+
+    return walks
diff --git a/cogdl/utils/rwalk/rwalk.c b/cogdl/utils/rwalk/rwalk.c
new file mode 100644
index 00000000..292ce618
--- /dev/null
+++ b/cogdl/utils/rwalk/rwalk.c
@@ -0,0 +1,39 @@
+#include "rwalk.h"
+#include <omp.h>
+#include <stdlib.h>
+
+
+void random_walk(int const* starts, int const* ptr, int const* neighs, int n, int num_walks,
+                 int num_steps, int seed, int nthread, float restart_prop, int* walks) {
+  if (nthread > 0) {
+    omp_set_num_threads(nthread);
+  }
+#pragma omp parallel
+  {
+    int thread_num = omp_get_thread_num();
+    unsigned int private_seed = (unsigned int)(seed + thread_num);
+#pragma omp for
+    for (int i = 0; i < n; i++) {
+      int offset, num_neighs;
+      for (int walk = 0; walk < num_walks; walk++) {
+        // int curr = i;
+        int curr = starts[i];
+        offset = i * num_walks * (num_steps + 1) + walk * (num_steps + 1);
+        walks[offset] = starts[i];
+        for (int step = 0; step < num_steps; step++) {
+          num_neighs = ptr[curr + 1] - ptr[curr];
+
+          if((restart_prop > 0) && (rand_r(&private_seed) / (double)RAND_MAX < restart_prop)){
+              curr = starts[i];
+          } else {
+            if (num_neighs > 0) {
+              curr = neighs[ptr[curr] + (rand_r(&private_seed) % num_neighs)];
+            }
+          }
+          walks[offset + step + 1] = curr;
+        }
+
+      }
+    }
+  }
+}
diff --git a/cogdl/utils/rwalk/rwalk.h b/cogdl/utils/rwalk/rwalk.h
new file mode 100644
index 00000000..8645b481
--- /dev/null
+++ b/cogdl/utils/rwalk/rwalk.h
@@ -0,0 +1,2 @@
+void random_walk(int const* starts, int const* ptr, int const* neighs, int n, int num_walks,
+                 int num_steps, int seed, int nthread, float restart_prop, int* walks);
\ No newline at end of file
diff --git a/cogdl/utils/sampling.py b/cogdl/utils/sampling.py
index a1db922c..d4460bb2 100644
--- a/cogdl/utils/sampling.py
+++ b/cogdl/utils/sampling.py
@@ -4,6 +4,8 @@
 import scipy.sparse as sp
 import random
 
+# from cogdl.utils.rwalk import random_walk as c_random_walk
+
 
 @numba.njit(cache=True, parallel=True)
 def random_walk(start, length, indptr, indices, p=0.0):
@@ -17,9 +19,9 @@ def random_walk(start, length, indptr, indices, p=0.0):
     Return:
         list(np.array(dtype=np.int32))
     """
-    result = [np.zeros(0, dtype=np.int32)] * len(start)
-    for node in start:
-        result[node] = _random_walk(node, length, indptr, indices, p)
+    result = [np.zeros(length, dtype=np.int32)] * len(start)
+    for i in numba.prange(len(start)):
+        result[i] = _random_walk(start[i], length, indptr, indices, p)
     return result
 
 
@@ -78,6 +80,8 @@ def walk(self, start, walk_length, restart_p=0.0):
         assert self.indptr is not None, "Please build the adj_list first"
         if isinstance(start, torch.Tensor):
             start = start.cpu().numpy()
+        if isinstance(start, list):
+            start = np.asarray(start, dtype=np.int32)
         result = random_walk(start, walk_length, self.indptr, self.indices, restart_p)
         result = np.array(result, dtype=np.int64)
         return result
diff --git a/cogdl/utils/spmm_utils.py b/cogdl/utils/spmm_utils.py
index 7c07a342..5d561ea6 100644
--- a/cogdl/utils/spmm_utils.py
+++ b/cogdl/utils/spmm_utils.py
@@ -94,29 +94,44 @@ def edge_softmax(graph, edge_val):
     Returns:
         Softmax values of edge values for nodes
     """
-    edge_val_max = edge_val.max().item()
-    while edge_val_max > 10:
-        edge_val -= edge_val / 2
+    csr_edge_softmax = CONFIGS["csr_edge_softmax"]
+    if csr_edge_softmax is not None and edge_val.device.type != "cpu":
+        edge_val = edge_val.view(-1, 1)
+        val = csr_edge_softmax(graph.row_indptr.int(), edge_val)
+        val = val.view(-1)
+        return val
+    else:
         edge_val_max = edge_val.max().item()
+        while edge_val_max > 10:
+            edge_val -= edge_val / 2
+            edge_val_max = edge_val.max().item()
 
-    with graph.local_graph():
-        edge_val = torch.exp(edge_val)
-        graph.edge_weight = edge_val
-        x = torch.ones(graph.num_nodes, 1).to(edge_val.device)
-        node_sum = spmm(graph, x).squeeze()
-        row = graph.edge_index[0]
-        softmax_values = edge_val / node_sum[row]
-        return softmax_values
+        with graph.local_graph():
+            edge_val = torch.exp(edge_val)
+            graph.edge_weight = edge_val
+            x = torch.ones(graph.num_nodes, 1).to(edge_val.device)
+            node_sum = spmm(graph, x).squeeze()
+            row = graph.edge_index[0]
+            softmax_values = edge_val / node_sum[row]
+            return softmax_values
 
 
 def mul_edge_softmax(graph, edge_val):
     """
+    Args:
+        graph: cogdl.Graph
+        edge_val: torch.Tensor, shape=(E, d)
     Returns:
         Softmax values of multi-dimension edge values. shape: [E, H]
     """
     csr_edge_softmax = CONFIGS["csr_edge_softmax"]
     if csr_edge_softmax is not None and edge_val.device.type != "cpu":
-        val = csr_edge_softmax(graph.row_indptr.int(), edge_val)
+        if len(edge_val.shape) == 1:
+            edge_val = edge_val.view(-1, 1)
+            val = csr_edge_softmax(graph.row_indptr.int(), edge_val)
+            val = val.view(-1)
+        else:
+            val = csr_edge_softmax(graph.row_indptr.int(), edge_val)
         return val
     else:
         val = []
diff --git a/cogdl/utils/transform.py b/cogdl/utils/transform.py
new file mode 100644
index 00000000..6b82a4ba
--- /dev/null
+++ b/cogdl/utils/transform.py
@@ -0,0 +1,87 @@
+from typing import Optional, Tuple
+
+import torch
+
+from cogdl.utils.graph_utils import symmetric_normalization, row_normalization
+
+
+class DropFeatures(torch.nn.Module):
+    def __init__(self, drop_rate):
+        super(DropFeatures, self).__init__()
+        self.drop_rate = drop_rate
+
+    def forward(self, x):
+        return dropout_features(x, self.drop_rate, training=self.training)
+
+
+class DropEdge(torch.nn.Module):
+    def __init__(self, drop_rate: float = 0.5, renorm: Optional[str] = "sym"):
+        super(DropEdge, self).__init__()
+        self.drop_rate = drop_rate
+        self.renorm = renorm
+
+    def forward(self, edge_index, edge_weight=None):
+        return dropout_adj(edge_index, edge_weight, self.drop_rate, self.renorm, self.training)
+
+
+class DropNode(torch.nn.Module):
+    def __init__(self, drop_rate=0.5):
+        super(DropNode, self).__init__()
+        self.drop_rate = drop_rate
+
+    def forward(self, x):
+        return drop_node(x, self.drop_rate, self.training)
+
+
+def filter_adj(row, col, edge_attr, mask):
+    return (row[mask], col[mask]), None if edge_attr is None else edge_attr[mask]
+
+
+def dropout_adj(
+    edge_index: Tuple,
+    edge_weight: Optional[torch.Tensor] = None,
+    drop_rate: float = 0.5,
+    renorm: Optional[str] = "sym",
+    training: bool = False,
+):
+    if not training or drop_rate == 0:
+        if edge_weight is None:
+            edge_weight = torch.ones(edge_index[0].shape[0], device=edge_index[0].device)
+        return edge_index, edge_weight
+
+    if drop_rate < 0.0 or drop_rate > 1.0:
+        raise ValueError("Dropout probability has to be between 0 and 1, " "but got {}".format(drop_rate))
+
+    row, col = edge_index
+    num_nodes = int(max(row.max(), col.max())) + 1
+    self_loop = row == col
+    mask = torch.full((row.shape[0],), 1 - drop_rate, dtype=torch.float, device=row.device)
+    mask = torch.bernoulli(mask).to(torch.bool)
+    mask = self_loop | mask
+    edge_index, edge_weight = filter_adj(row, col, edge_weight, mask)
+    if renorm == "sym":
+        edge_weight = symmetric_normalization(num_nodes, edge_index[0], edge_index[1])
+    elif renorm == "row":
+        edge_weight = row_normalization(num_nodes, edge_index[0], edge_index[1])
+    return edge_index, edge_weight
+
+
+def dropout_features(x: torch.Tensor, droprate: float, training=True):
+    n = x.shape[1]
+    drop_rates = torch.ones(n, device=x.device) * droprate
+    if training:
+        masks = torch.bernoulli(1.0 - drop_rates).view(1, -1).expand_as(x)
+        masks = masks.to(x.device)
+        masks = masks.to(x.device)
+        x = masks * x
+    return x
+
+
+def drop_node(x, drop_rate=0.5, training=True):
+    n = x.shape[0]
+    drop_rates = torch.ones(n) * drop_rate
+    if training:
+        masks = torch.bernoulli(1.0 - drop_rates).unsqueeze(1)
+        x = masks.to(x.device) * x
+        x = x / drop_rate
+    return x
diff --git a/cogdl/utils/utils.py b/cogdl/utils/utils.py
index 96a12818..73ab7a50 100644
--- a/cogdl/utils/utils.py
+++ b/cogdl/utils/utils.py
@@ -9,6 +9,7 @@
 
 import numpy as np
 import torch
+import torch.nn as nn
 import torch.nn.functional as F
 from tabulate import tabulate
 
@@ -147,15 +148,15 @@ def identity_act(input, inplace=True):
 
 def get_activation(act: str):
     if act == "relu":
-        return F.relu
+        return nn.ReLU()
     elif act == "sigmoid":
-        return torch.sigmoid
+        return nn.Sigmoid()
     elif act == "tanh":
-        return torch.tanh
+        return nn.Tanh()
     elif act == "gelu":
-        return F.gelu
+        return nn.GELU()
     elif act == "prelu":
-        return F.prelu
+        return nn.PReLU()
     elif act == "identity":
         return identity_act
     else:
@@ -209,6 +210,7 @@ def batch_max_pooling(x, batch):
 
 def tabulate_results(results_dict):
     # Average for different seeds
+    # {"model1_dataset": [dict(acc=1), dict(acc=2)], "model2_dataset": [dict(acc=1),...]}
     tab_data = []
     for variant in results_dict:
         results = np.array([list(res.values()) for res in results_dict[variant]])
diff --git a/cogdl/wrappers/__init__.py b/cogdl/wrappers/__init__.py
new file mode 100644
index 00000000..0255299d
--- /dev/null
+++ b/cogdl/wrappers/__init__.py
@@ -0,0 +1,8 @@
+from .data_wrapper import try_import_data_wrapper, register_data_wrapper, fetch_data_wrapper
+from .model_wrapper import (
+    try_import_model_wrapper,
+    register_model_wrapper,
+    fetch_model_wrapper,
+    ModelWrapper,
+    EmbeddingModelWrapper,
+)
diff --git a/cogdl/wrappers/data_wrapper/__init__.py b/cogdl/wrappers/data_wrapper/__init__.py
new file mode 100644
index 00000000..e1604eef
--- /dev/null
+++ b/cogdl/wrappers/data_wrapper/__init__.py
@@ -0,0 +1,61 @@
+from .base_data_wrapper import DataWrapper
+import os
+import importlib
+
+
+DATAMODULE_REGISTRY = {}
+SUPPORTED_DATAMODULE = {}
+
+
+def register_data_wrapper(name):
+    """
+    New data wrapper types can be added to cogdl with the :func:`register_data_wrapper`
+    function decorator.
+
+    Args:
+        name (str): the name of the data_wrapper
+    """
+
+    def register_data_wrapper_cls(cls):
+        if name in DATAMODULE_REGISTRY:
+            raise ValueError("Cannot register duplicate data_wrapper ({})".format(name))
+        if not issubclass(cls, DataWrapper):
+            raise ValueError("({}: {}) must extend DataWrapper".format(name, cls.__name__))
+        DATAMODULE_REGISTRY[name] = cls
+        cls.model_name = name
+        return cls
+
+    return register_data_wrapper_cls
+
+
+def scan_data_wrappers():
+    global SUPPORTED_DATAMODULE
+    dirname = os.path.dirname(__file__)
+    dir_names = [x for x in os.listdir(dirname) if not x.startswith("__")]
+    dirs = [os.path.join(dirname, x) for x in dir_names]
+    dirs_names = [(x, y) for x, y in zip(dirs, dir_names) if os.path.isdir(x)]
+    dw_dict = SUPPORTED_DATAMODULE
+    for _dir, _name in dirs_names:
+        files = os.listdir(_dir)
+        dw = [x.split(".")[0] for x in files]
+        dw = [x for x in dw if not x.startswith("__")]
+        path = [f"cogdl.wrappers.data_wrapper.{_name}.{x}" for x in dw]
+        for x, y in zip(dw, path):
+            dw_dict[x] = y
+
+
+def try_import_data_wrapper(name):
+    if name in DATAMODULE_REGISTRY:
+        return
+    if name in SUPPORTED_DATAMODULE:
+        importlib.import_module(SUPPORTED_DATAMODULE[name])
+    else:
+        raise NotImplementedError(f"`{name}` data_wrapper is not implemented.")
+
+
+def fetch_data_wrapper(name):
+    try_import_data_wrapper(name)
+    return DATAMODULE_REGISTRY[name]
+
+
+scan_data_wrappers()
diff --git a/cogdl/wrappers/data_wrapper/base_data_wrapper.py b/cogdl/wrappers/data_wrapper/base_data_wrapper.py
new file mode 100644
index 00000000..d858c1d1
--- /dev/null
+++ b/cogdl/wrappers/data_wrapper/base_data_wrapper.py
@@ -0,0 +1,273 @@
+import numpy as np
+import torch
+from torch.utils.data import DataLoader
+from cogdl.data import Graph
+
+
+class DataWrapper(object):
+    @staticmethod
+    def add_args(parser):
+        pass
+
+    def __init__(self, dataset=None):
+        if dataset is not None:
+            if hasattr(dataset, "get_loss_fn"):
+                self.__loss_fn__ = dataset.get_loss_fn()
+            if hasattr(dataset, "get_evaluator"):
+                self.__evaluator__ = dataset.get_evaluator()
+        else:
+            self.__loss_fn__ = None
+            self.__evaluator__ = None
+        self.__dataset__ = dataset
+        self.__training_data, self.__val_data, self.__test_data = None, None, None
+        self.__num_training_data, self.__num_val_data, self.__num_test_data = 0, 0, 0
+        self.__prepare_dataloader_per_epoch__ = False
+        self.__back_to_cpu__ = False
+
+    @property
+    def data_back_to_cpu(self):
+        return (
+            isinstance(self.__val_data.raw_data, Graph) or isinstance(self.__test_data.raw_data, Graph)
+        ) and not isinstance(self.__training_data.raw_data, Graph)
+
+    def get_train_dataset(self):
+        """
+        Return the `wrapped` dataset for specific usage.
+        For example, return `ClusteredDataset` in cluster_dw for DDP training.
+        """
+        raise NotImplementedError
+
+    def get_val_dataset(self):
+        """
+        Similar to `self.get_train_dataset` but for validation.
+        """
+        raise NotImplementedError
+
+    def get_test_dataset(self):
+        """
+        Similar to `self.get_train_dataset` but for test.
+        """
+        raise NotImplementedError
+
+    def train_wrapper(self):
+        """
+        Return:
+            1. DataLoader
+            2. cogdl.Graph
+            3. list of DataLoader or Graph
+        Any other data formats other than DataLoader will not be traversed
+        """
+        pass
+
+    def val_wrapper(self):
+        pass
+
+    def test_wrapper(self):
+        pass
+
+    def evaluation_wrapper(self):
+        if self.__dataset__ is None:
+            self.__dataset__ = getattr(self, "dataset", None)
+        if self.__dataset__ is not None:
+            return self.__dataset__
+
+    def train_transform(self, batch):
+        return batch
+
+    def val_transform(self, batch):
+        return batch
+
+    def test_transform(self, batch):
+        return batch
+
+    def pre_transform(self):
+        """Data Preprocessing before all runs"""
+        pass
+
+    def pre_stage(self, stage, model_w_out):
+        """Processing before each run"""
+        pass
+
+    def post_stage(self, stage, model_w_out):
+        """Processing after each run"""
+        pass
+
+    def refresh_per_epoch(self, name="train"):
+        self.__prepare_dataloader_per_epoch__ = True
+
+    def __refresh_per_epoch__(self):
+        return self.__prepare_dataloader_per_epoch__
+
+    def get_default_loss_fn(self):
+        return self.__loss_fn__
+
+    def get_default_evaluator(self):
+        return self.__evaluator__
+
+    def get_dataset(self):
+        if self.__dataset__ is None:
+            self.__dataset__ = getattr(self, "dataset", None)
+        return self.__dataset__
+
+    def prepare_training_data(self):
+        train_data = self.train_wrapper()
+        if train_data is not None:
+            self.__training_data = OnLoadingWrapper(train_data, self.train_transform)
+
+    def prepare_val_data(self):
+        val_data = self.val_wrapper()
+        if val_data is not None:
+            self.__val_data = OnLoadingWrapper(val_data, self.val_transform)
+
+    def prepare_test_data(self):
+        test_data = self.test_wrapper()
+        if test_data is not None:
+            self.__test_data = OnLoadingWrapper(test_data, self.test_transform)
+
+    def set_train_data(self, x):
+        self.__training_data = x
+
+    def set_val_data(self, x):
+        self.__val_data = x
+
+    def set_test_data(self, x):
+        self.__test_data = x
+
+    def on_train_wrapper(self):
+        if self.__training_data is None:
+            return None
+
+        if self.__prepare_dataloader_per_epoch__:
+            # TODO: reserve parameters for `prepare training data`
+            self.prepare_training_data()
+        return self.__training_data
+
+    def on_val_wrapper(self):
+        return self.__val_data
+
+    def on_test_wrapper(self):
+        return self.__test_data
+
+    def train(self):
+        if self.__dataset__ is None:
+            self.__dataset__ = getattr(self, "dataset", None)
+        if self.__dataset__ is not None and isinstance(self.__dataset__.data, Graph):
+            self.__dataset__.data.train()
+
+    def eval(self):
+        if self.__dataset__ is None:
+            self.__dataset__ = getattr(self, "dataset", None)
+        if self.__dataset__ is not None and isinstance(self.__dataset__.data, Graph):
+            self.__dataset__.data.eval()
+
+
+class OnLoadingWrapper(object):
+    def __init__(self, data, transform):
+        """
+        Args:
+            data: `data` or `dataset`, that it, `cogdl.Graph` or `DataLoader`
+        """
+        self.raw_data = data
+        self.data = self.__process_iterative_data__(data)
+        self.__num_training_data = self.__get_min_len__(self.data)
+        self.wrapped_data = self.__wrap_iteration__(self.data)
+        self.ptr = 0
+        self.transform = transform
+
+    def __next__(self):
+        if self.ptr < self.__num_training_data:
+            self.ptr += 1
+            batch = self.__next_batch__(self.wrapped_data)
+            return self.transform(batch)
+        else:
+            self.ptr = 0
+            # re-wrap the dataset per epoch
+            self.wrapped_data = self.__wrap_iteration__(self.data)
+            raise StopIteration
+
+    def __iter__(self):
+        return self
+
+    def __len__(self):
+        return self.__num_training_data
+
+    def get_dataset_from_loader(self):
+        return self.raw_data
+
+    def __wrap_iteration__(self, inputs):
+        # if isinstance(inputs, tuple):
+        #     inputs = list(inputs)
+        def iter_func(in_x):
+            if isinstance(in_x, list) or isinstance(in_x, DataLoader):
+                for item in in_x:
+                    yield item
+            else:
+                yield in_x
+
+        if isinstance(inputs, list):
+            outputs = [None] * len(inputs)
+            for i, item in enumerate(inputs):
+                outputs[i] = self.__wrap_iteration__(item)
+        elif isinstance(inputs, dict):
+            outputs = {key: None for key in inputs.keys()}
+            for key, val in inputs.items():
+                outputs[key] = self.__wrap_iteration__(val)
+        else:
+            # return LoaderWrapper(inputs)
+            return iter_func(inputs)
+        return outputs
+
+    def __process_iterative_data__(self, inputs):
+        if inputs is None:
+            return None
+        # if isinstance(inputs, tuple):
+        #     inputs = list(inputs)
+
+        if isinstance(inputs, list):
+            for i, item in enumerate(inputs):
+                inputs[i] = self.__process_iterative_data__(item)
+        elif isinstance(inputs, dict):
+            for key, val in inputs.items():
+                inputs[key] = self.__process_iterative_data__(val)
+        else:
+            # return self.__batch_wrapper__(inputs)
+            return inputs
+        return inputs
+
+    def __next_batch__(self, inputs):
+        # if isinstance(inputs, tuple):
+        #     inputs = list(inputs)
+
+        if isinstance(inputs, list):
+            outputs = [None] * len(inputs)
+            for i, item in enumerate(inputs):
+                outputs[i] = self.__next_batch__(item)
+        elif isinstance(inputs, dict):
+            outputs = {key: None for key in inputs.keys()}
+            for key, val in inputs.items():
+                outputs[key] = self.__next_batch__(val)
+        else:
+            return next(inputs)
+        return outputs
+
+    def __get_min_len__(self, inputs):
+        if inputs is None:
+            return None
+
+        # if isinstance(inputs, tuple):
+        #     inputs = list(inputs)
+        if isinstance(inputs, list):
+            outputs = [0] * len(inputs)
+            for i, item in enumerate(inputs):
+                inputs[i] = self.__get_min_len__(item)
+            return np.min(outputs)
+        # elif isinstance(inputs, dict):
+        #     outputs = {key: 0 for key in inputs.keys()}
+        #     for i, val in enumerate(inputs.values()):
+        #         outputs[i] = self.__get_min_len__(val)
+        #     return np.min(list(outputs.values()))
+        else:
+            if isinstance(inputs, DataLoader):
+                return len(inputs)
+            else:
+                return 1
diff --git a/cogdl/wrappers/data_wrapper/graph_classification/__init__.py b/cogdl/wrappers/data_wrapper/graph_classification/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/cogdl/wrappers/data_wrapper/graph_classification/graph_classification_dw.py b/cogdl/wrappers/data_wrapper/graph_classification/graph_classification_dw.py
new file mode 100644
index 00000000..a4ff3154
--- /dev/null
+++ b/cogdl/wrappers/data_wrapper/graph_classification/graph_classification_dw.py
@@ -0,0 +1,72 @@
+import torch
+
+from .. import register_data_wrapper, DataWrapper
+from cogdl.wrappers.tools.wrapper_utils import node_degree_as_feature, split_dataset
+from cogdl.data import DataLoader
+from cogdl.models.nn.patchy_san import get_single_feature
+
+
+@register_data_wrapper("graph_classification_dw")
+class GraphClassificationDataWrapper(DataWrapper):
+    @staticmethod
+    def add_args(parser):
+        # fmt: off
+        parser.add_argument("--degree-node-features", action="store_true",
+                            help="Use one-hot degree vector as input node features")
+        # parser.add_argument("--kfold", action="store_true", help="Use 10-fold cross-validation")
+        parser.add_argument("--train-ratio", type=float, default=0.5)
+        parser.add_argument("--test-ratio", type=float, default=0.3)
+        parser.add_argument("--batch-size", type=int, default=16)
+        # fmt: on
+
+    def __init__(self, dataset, degree_node_features=False, batch_size=32, train_ratio=0.5, test_ratio=0.3):
+        super(GraphClassificationDataWrapper, self).__init__(dataset)
+        self.dataset = dataset
+        self.degree_node_features = degree_node_features
+        self.train_ratio = train_ratio
+        self.test_ratio = test_ratio
+        self.batch_size = batch_size
+        self.split_idx = None
+
+    def train_wrapper(self):
+        return DataLoader(self.dataset[self.split_idx[0]], batch_size=self.batch_size, shuffle=True, num_workers=4)
+
+    def val_wrapper(self):
+        if self.split_idx[1] is not None:
+            return DataLoader(self.dataset[self.split_idx[1]], batch_size=self.batch_size, shuffle=False, num_workers=4)
+
+    def test_wrapper(self):
+        return DataLoader(self.dataset[self.split_idx[2]], batch_size=self.batch_size, shuffle=False, num_workers=4)
+
+    def pre_transform(self):
+        if self.degree_node_features and self.dataset.data.x is None:
+            self.dataset = node_degree_as_feature(self.dataset)
+        train_idx, val_idx, test_idx = split_dataset(len(self.dataset), self.train_ratio, self.test_ratio)
+        self.split_idx = [train_idx, val_idx, test_idx]
+
+
+@register_data_wrapper("patchy_san_dw")
+class PATCHY_SAN_DataWrapper(GraphClassificationDataWrapper):
+    @staticmethod
+    def add_args(parser):
+        GraphClassificationDataWrapper.add_args(parser)
+        parser.add_argument("--num-sample", default=30, type=int, help="Number of chosen vertexes")
+        parser.add_argument("--num-neighbor", default=10, type=int, help="Number of neighbor in constructing features")
+
+    def __init__(self, dataset, num_sample, num_neighbor, *args, **kwargs):
+        super(PATCHY_SAN_DataWrapper, self).__init__(dataset, *args, **kwargs)
+        self.sample = num_sample
+        self.dataset = dataset
+        self.neighbor = num_neighbor
+
+    def pre_transform(self):
+        dataset = self.dataset
+        num_features = dataset.data.num_features
+        num_classes = dataset.data.num_classes
+        for i, data in enumerate(dataset):
+            new_feature = get_single_feature(
+                dataset[i], num_features, num_classes, self.sample, self.neighbor, self.stride
+            )
+            dataset[i].x = torch.from_numpy(new_feature)
+        self.dataset = dataset
+        super(PATCHY_SAN_DataWrapper, self).pre_transform()
diff --git a/cogdl/wrappers/data_wrapper/graph_classification/graph_embedding_dw.py b/cogdl/wrappers/data_wrapper/graph_classification/graph_embedding_dw.py
new file mode 100644
index 00000000..4a930898
--- /dev/null
+++ b/cogdl/wrappers/data_wrapper/graph_classification/graph_embedding_dw.py
@@ -0,0 +1,32 @@
+import numpy as np
+
+from .. import register_data_wrapper, DataWrapper
+from cogdl.wrappers.tools.wrapper_utils import node_degree_as_feature
+
+
+@register_data_wrapper("graph_embedding_dw")
+class GraphEmbeddingDataWrapper(DataWrapper):
+    @staticmethod
+    def add_args(parser):
+        # fmt: off
+        parser.add_argument("--degree-node-features", action="store_true",
+                            help="Use one-hot degree vector as input node features")
+        # fmt: on
+
+    def __init__(self, dataset, degree_node_features=False):
+        super(GraphEmbeddingDataWrapper, self).__init__(dataset)
+        self.dataset = dataset
+        self.degree_node_features = degree_node_features
+
+    def train_wrapper(self):
+        return self.dataset
+
+    def test_wrapper(self):
+        if self.dataset[0].y.shape[0] > 1:
+            return np.array([g.y.numpy() for g in self.dataset])
+        else:
+            return np.array([g.y.item() for g in self.dataset])
+
+    def pre_transform(self):
+        if self.degree_node_features:
+            self.dataset = node_degree_as_feature(self.dataset)
diff --git a/cogdl/wrappers/data_wrapper/graph_classification/infograph_dw.py b/cogdl/wrappers/data_wrapper/graph_classification/infograph_dw.py
new file mode 100644
index 00000000..23d498a6
--- /dev/null
+++ b/cogdl/wrappers/data_wrapper/graph_classification/infograph_dw.py
@@ -0,0 +1,8 @@
+from .. import register_data_wrapper
+from .graph_classification_dw import GraphClassificationDataWrapper
+
+
+@register_data_wrapper("infograph_dw")
+class InfoGraphDataWrapper(GraphClassificationDataWrapper):
+    def test_wrapper(self):
+        return self.dataset
diff --git a/cogdl/wrappers/data_wrapper/heterogeneous/__init__.py b/cogdl/wrappers/data_wrapper/heterogeneous/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/cogdl/wrappers/data_wrapper/heterogeneous/heterogeneous_embedding_dw.py b/cogdl/wrappers/data_wrapper/heterogeneous/heterogeneous_embedding_dw.py
new file mode 100644
index 00000000..3d8df67d
--- /dev/null
+++ b/cogdl/wrappers/data_wrapper/heterogeneous/heterogeneous_embedding_dw.py
@@ -0,0 +1,15 @@
+from .. import register_data_wrapper, DataWrapper
+
+
+@register_data_wrapper("heterogeneous_embedding_dw")
+class HeterogeneousEmbeddingDataWrapper(DataWrapper):
+    def __init__(self, dataset):
+        super(HeterogeneousEmbeddingDataWrapper, self).__init__()
+
+        self.dataset = dataset
+
+    def train_wrapper(self):
+        return self.dataset.data
+
+    def test_wrapper(self):
+        return self.dataset.data
diff --git a/cogdl/wrappers/data_wrapper/heterogeneous/heterogeneous_gnn_dw.py b/cogdl/wrappers/data_wrapper/heterogeneous/heterogeneous_gnn_dw.py
new file mode 100644
index 00000000..9d1025b8
--- /dev/null
+++ b/cogdl/wrappers/data_wrapper/heterogeneous/heterogeneous_gnn_dw.py
@@ -0,0 +1,18 @@
+from .. import register_data_wrapper, DataWrapper
+
+
+@register_data_wrapper("heterogeneous_gnn_dw")
+class HeterogeneousGNNDataWrapper(DataWrapper):
+    def __init__(self, dataset):
+        super(HeterogeneousGNNDataWrapper, self).__init__(dataset=dataset)
+
+        self.dataset = dataset
+
+    def train_wrapper(self):
+        return self.dataset
+
+    def val_wrapper(self):
+        return self.dataset
+
+    def test_wrapper(self):
+        return self.dataset
diff --git a/cogdl/wrappers/data_wrapper/heterogeneous/multiplex_embedding_dw.py b/cogdl/wrappers/data_wrapper/heterogeneous/multiplex_embedding_dw.py
new file mode 100644
index 00000000..1b3e5f95
--- /dev/null
+++ b/cogdl/wrappers/data_wrapper/heterogeneous/multiplex_embedding_dw.py
@@ -0,0 +1,15 @@
+from .. import register_data_wrapper, DataWrapper
+
+
+@register_data_wrapper("multiplex_embedding_dw")
+class MultiplexEmbeddingDataWrapper(DataWrapper):
+    def __init__(self, dataset):
+        super(MultiplexEmbeddingDataWrapper, self).__init__()
+
+        self.dataset = dataset
+
+    def train_wrapper(self):
+        return self.dataset.data.train_data
+
+    def test_wrapper(self):
+        return self.dataset.data.test_data
diff --git a/cogdl/wrappers/data_wrapper/link_predicttion/__init__.py b/cogdl/wrappers/data_wrapper/link_predicttion/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/cogdl/wrappers/data_wrapper/link_predicttion/embedding_link_prediction_dw.py b/cogdl/wrappers/data_wrapper/link_predicttion/embedding_link_prediction_dw.py
new file mode 100644
index 00000000..e500bd64
--- /dev/null
+++ b/cogdl/wrappers/data_wrapper/link_predicttion/embedding_link_prediction_dw.py
@@ -0,0 +1,91 @@
+import random
+import networkx as nx
+import numpy as np
+import torch
+
+from .. import DataWrapper, register_data_wrapper
+from cogdl.data import Graph
+
+
+@register_data_wrapper("embedding_link_prediction_dw")
+class EmbeddingLinkPredictionDataWrapper(DataWrapper):
+    @staticmethod
+    def add_args(parser):
+        # fmt: off
+        parser.add_argument("--negative-ratio", type=int, default=5)
+        # fmt: on
+
+    def __init__(self, dataset, negative_ratio):
+        super(EmbeddingLinkPredictionDataWrapper, self).__init__(dataset)
+        self.dataset = dataset
+        self.negative_ratio = negative_ratio
+        self.train_data, self.test_data = None, None
+
+    def train_wrapper(self):
+        return self.train_data
+
+    def test_wrapper(self):
+        return self.test_data
+
+    def pre_transform(self):
+        row, col = self.dataset.data.edge_index
+        edge_list = list(zip(row.numpy(), col.numpy()))
+        edge_set = set()
+        for edge in edge_list:
+            if (edge[0], edge[1]) not in edge_set and (edge[1], edge[0]) not in edge_set:
+                edge_set.add(edge)
+        edge_list = list(edge_set)
+        train_edges, test_edges = divide_data(edge_list, [0.90, 0.10])
+        self.test_data = gen_node_pairs(train_edges, test_edges, self.negative_ratio)
+        train_edges = np.array(train_edges).transpose()
+        train_edges = torch.from_numpy(train_edges)
+        self.train_data = Graph(edge_index=train_edges)
+
+
+def divide_data(input_list, division_rate):
+    local_division = len(input_list) * np.cumsum(np.array(division_rate))
+    random.shuffle(input_list)
+    return [
+        input_list[int(round(local_division[i - 1])) if i > 0 else 0 : int(round(local_division[i]))]
+        for i in range(len(local_division))
+    ]
+
+
+def randomly_choose_false_edges(nodes, true_edges, num):
+    true_edges_set = set(true_edges)
+    tmp_list = list()
+    all_flag = False
+    for _ in range(num):
+        trial = 0
+        while True:
+            x = nodes[random.randint(0, len(nodes) - 1)]
+            y = nodes[random.randint(0, len(nodes) - 1)]
+            trial += 1
+            if trial >= 1000:
+                all_flag = True
+                break
+            if x != y and (x, y) not in true_edges_set and (y, x) not in true_edges_set:
+                tmp_list.append((x, y))
+                break
+        if all_flag:
+            break
+    return tmp_list
+
+
+def gen_node_pairs(train_data, test_data, negative_ratio=5):
+    G = nx.Graph()
+    G.add_edges_from(train_data)
+
+    training_nodes = set(list(G.nodes()))
+    test_true_data = []
+    for u, v in test_data:
+        if u in training_nodes and v in training_nodes:
+            test_true_data.append((u, v))
+    test_false_data = randomly_choose_false_edges(list(training_nodes), train_data, len(test_data) * negative_ratio)
+    return (test_true_data, test_false_data)
+
+
+def get_score(embs, node1, node2):
+    vector1 = embs[int(node1)]
+    vector2 = embs[int(node2)]
+    return np.dot(vector1, vector2) / (np.linalg.norm(vector1) * np.linalg.norm(vector2))
diff --git a/cogdl/wrappers/data_wrapper/link_predicttion/gnn_kg_link_prediction_dw.py b/cogdl/wrappers/data_wrapper/link_predicttion/gnn_kg_link_prediction_dw.py
new file mode 100644
index 00000000..77c2aa1d
--- /dev/null
+++ b/cogdl/wrappers/data_wrapper/link_predicttion/gnn_kg_link_prediction_dw.py
@@ -0,0 +1,18 @@
+from .. import register_data_wrapper, DataWrapper
+
+
+@register_data_wrapper("gnn_kg_link_prediction_dw")
+class GNNKGLinkPredictionModelWrapper(DataWrapper):
+    def __init__(self, dataset):
+        super(GNNKGLinkPredictionModelWrapper, self).__init__(dataset)
+        self.dataset = dataset
+        self.edge_set = None
+
+    def train_wrapper(self):
+        return self.dataset.data
+
+    def val_wrapper(self):
+        return self.dataset.data
+
+    def test_wrapper(self):
+        return self.dataset.data
diff --git a/cogdl/wrappers/data_wrapper/link_predicttion/gnn_link_prediction_dw.py b/cogdl/wrappers/data_wrapper/link_predicttion/gnn_link_prediction_dw.py
new file mode 100644
index 00000000..2f74bd9d
--- /dev/null
+++ b/cogdl/wrappers/data_wrapper/link_predicttion/gnn_link_prediction_dw.py
@@ -0,0 +1,84 @@
+import numpy as np
+import torch
+
+from .. import DataWrapper, register_data_wrapper
+
+
+@register_data_wrapper("gnn_link_prediction_dw")
+class GNNLinkPredictionDataWrapper(DataWrapper):
+    def __init__(self, dataset):
+        super(GNNLinkPredictionDataWrapper, self).__init__(dataset)
+        self.dataset = dataset
+
+    def train_wrapper(self):
+        return self.dataset.data
+
+    def val_wrapper(self):
+        return self.dataset.data
+
+    def test_wrapper(self):
+        return self.dataset.data
+
+    def pre_transform(self):
+        data = self.dataset.data
+        num_nodes = data.x.shape[0]
+        (
+            (train_edges, val_edges, test_edges),
+            (val_false_edges, test_false_edges),
+        ) = self.train_test_edge_split(data.edge_index, num_nodes)
+        data.train_edges = train_edges
+        data.val_edges = val_edges
+        data.test_edges = test_edges
+        data.val_neg_edges = val_false_edges
+        data.test_neg_edges = test_false_edges
+
+    @staticmethod
+    def train_test_edge_split(edge_index, num_nodes, val_ratio=0.1, test_ratio=0.2):
+        row, col = edge_index
+        mask = row > col
+        row, col = row[mask], col[mask]
+        num_edges = row.size(0)
+
+        perm = torch.randperm(num_edges)
+        row, col = row[perm], col[perm]
+
+        num_val = int(num_edges * val_ratio)
+        num_test = int(num_edges * test_ratio)
+
+        index = [[0, num_val], [num_val, num_val + num_test], [num_val + num_test, -1]]
+        sampled_rows = [row[l:r] for l, r in index]  # noqa E741
+        sampled_cols = [col[l:r] for l, r in index]  # noqa E741
+
+        # sample false edges
+        num_false = num_val + num_test
+        row_false = np.random.randint(0, num_nodes, num_edges * 5)
+        col_false = np.random.randint(0, num_nodes, num_edges * 5)
+
+        indices_false = row_false * num_nodes + col_false
+        indices_true = row.cpu().numpy() * num_nodes + col.cpu().numpy()
+        indices_false = list(set(indices_false).difference(indices_true))
+        indices_false = np.array(indices_false)
+        row_false = indices_false // num_nodes
+        col_false = indices_false % num_nodes
+
+        mask = row_false > col_false
+        row_false = row_false[mask]
+        col_false = col_false[mask]
+
+        edge_index_false = np.stack([row_false, col_false])
+        if edge_index[0].shape[0] < num_false:
+            ratio = edge_index_false.shape[1] / num_false
+            num_val = int(ratio * num_val)
+            num_test = int(ratio * num_test)
+        val_false_edges = torch.from_numpy(edge_index_false[:, 0:num_val])
+        test_fal_edges = torch.from_numpy(edge_index_false[:, num_val : num_test + num_val])
+
+        def to_undirected(_row, _col):
+            _edge_index = torch.stack([_row, _col], dim=0)
+            _r_edge_index = torch.stack([_col, _row], dim=0)
+            return torch.cat([_edge_index, _r_edge_index], dim=1)
+
+        train_edges = to_undirected(sampled_rows[2], sampled_cols[2])
+        val_edges = torch.stack([sampled_rows[0], sampled_cols[0]])
+        test_edges = torch.stack([sampled_rows[1], sampled_cols[1]])
+        return (train_edges, val_edges, test_edges), (val_false_edges, test_fal_edges)
diff --git a/cogdl/wrappers/data_wrapper/node_classification/__init__.py b/cogdl/wrappers/data_wrapper/node_classification/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/cogdl/wrappers/data_wrapper/node_classification/cluster_dw.py b/cogdl/wrappers/data_wrapper/node_classification/cluster_dw.py
new file mode 100644
index 00000000..58c85fcd
--- /dev/null
+++ b/cogdl/wrappers/data_wrapper/node_classification/cluster_dw.py
@@ -0,0 +1,44 @@
+from .. import DataWrapper, register_data_wrapper
+from cogdl.data.sampler import ClusteredLoader, ClusteredDataset
+
+
+@register_data_wrapper("cluster_dw")
+class ClusterWrapper(DataWrapper):
+    @staticmethod
+    def add_args(parser):
+        # fmt: off
+        parser.add_argument("--batch-size", type=int, default=20)
+        parser.add_argument("--n-cluster", type=int, default=100)
+        parser.add_argument("--method", type=str, default="metis")
+        # fmt: on
+
+    def __init__(self, dataset, method="metis", batch_size=20, n_cluster=100):
+        super(ClusterWrapper, self).__init__(dataset)
+        self.dataset = dataset
+        self.cluster_dataset = ClusteredDataset(dataset, n_cluster=n_cluster, batch_size=batch_size)
+        self.batch_size = batch_size
+        self.n_cluster = n_cluster
+        self.method = method
+
+    def train_wrapper(self):
+        self.dataset.data.train()
+        return ClusteredLoader(
+            self.cluster_dataset,
+            method=self.method,
+            batch_size=self.batch_size,
+            shuffle=True,
+            n_cluster=self.n_cluster,
+            # persistent_workers=True,
+            num_workers=0,
+        )
+
+    def get_train_dataset(self):
+        return self.cluster_dataset
+
+    def val_wrapper(self):
+        self.dataset.data.eval()
+        return self.dataset.data
+
+    def test_wrapper(self):
+        self.dataset.data.eval()
+        return self.dataset.data
diff --git a/cogdl/wrappers/data_wrapper/node_classification/graphsage_dw.py b/cogdl/wrappers/data_wrapper/node_classification/graphsage_dw.py
new file mode 100644
index 00000000..4fabaadd
--- /dev/null
+++ b/cogdl/wrappers/data_wrapper/node_classification/graphsage_dw.py
@@ -0,0 +1,80 @@
+from .. import DataWrapper, register_data_wrapper
+from cogdl.data.sampler import NeighborSampler, NeighborSamplerDataset
+
+
+@register_data_wrapper("graphsage_dw")
+class SAGEDataWrapper(DataWrapper):
+    @staticmethod
+    def add_args(parser):
+        # fmt: off
+        parser.add_argument("--batch-size", type=int, default=128)
+        parser.add_argument("--sample-size", type=int, nargs='+', default=[10, 10])
+        # fmt: on
+
+    def __init__(self, dataset, batch_size: int, sample_size: list):
+        super(SAGEDataWrapper, self).__init__(dataset)
+        self.dataset = dataset
+        self.train_dataset = NeighborSamplerDataset(
+            dataset, sizes=sample_size, batch_size=batch_size, mask=dataset.data.train_mask
+        )
+        self.val_dataset = NeighborSamplerDataset(
+            dataset, sizes=sample_size, batch_size=batch_size * 2, mask=dataset.data.val_mask
+        )
+        self.test_dataset = NeighborSamplerDataset(
+            dataset=self.dataset,
+            mask=None,
+            sizes=[-1],
+            batch_size=self.batch_size * 2,
+        )
+        self.x = self.dataset.data.x
+        self.y = self.dataset.data.y
+        self.batch_size = batch_size
+        self.sample_size = sample_size
+
+    def train_wrapper(self):
+        self.dataset.data.train()
+        return NeighborSampler(
+            dataset=self.train_dataset,
+            mask=self.dataset.data.train_mask,
+            sizes=self.sample_size,
+            num_workers=4,
+            shuffle=False,
+            batch_size=self.batch_size,
+        )
+
+    def val_wrapper(self):
+        self.dataset.data.eval()
+
+        return NeighborSampler(
+            dataset=self.val_dataset,
+            mask=self.dataset.data.val_mask,
+            sizes=self.sample_size,
+            batch_size=self.batch_size * 2,
+            shuffle=False,
+            num_workers=4,
+        )
+
+    def test_wrapper(self):
+        return (
+            self.dataset,
+            NeighborSampler(
+                dataset=self.test_dataset,
+                mask=None,
+                sizes=[-1],
+                batch_size=self.batch_size * 2,
+                shuffle=False,
+                num_workers=4,
+            ),
+        )
+
+    def train_transform(self, batch):
+        target_id, n_id, adjs = batch
+        x_src = self.x[n_id]
+        y = self.y[target_id]
+        return x_src, y, adjs
+
+    def val_transform(self, batch):
+        return self.train_transform(batch)
+
+    def get_train_dataset(self):
+        return self.train_dataset
diff --git a/cogdl/wrappers/data_wrapper/node_classification/m3s_dw.py b/cogdl/wrappers/data_wrapper/node_classification/m3s_dw.py
new file mode 100644
index 00000000..fa1a608d
--- /dev/null
+++ b/cogdl/wrappers/data_wrapper/node_classification/m3s_dw.py
@@ -0,0 +1,98 @@
+import numpy as np
+import scipy.sparse as sp
+import scipy.sparse.linalg as slinalg
+
+import torch
+
+from .. import register_data_wrapper
+from .node_classification_dw import FullBatchNodeClfDataWrapper
+
+
+@register_data_wrapper("m3s_dw")
+class M3SDataWrapper(FullBatchNodeClfDataWrapper):
+    @staticmethod
+    def add_args(parser):
+        # fmt: off
+        parser.add_argument("--label-rate", type=float, default=0.2)
+        parser.add_argument("--approximate", action="store_true")
+        parser.add_argument("--alpha", type=float, default=0.2)
+        # fmt: on
+
+    def __init__(self, dataset, label_rate, approximate, alpha):
+        super(M3SDataWrapper, self).__init__(dataset)
+        self.dataset = dataset
+        self.label_rate = label_rate
+        self.approximate = approximate
+        self.alpha = alpha
+
+    def pre_transform(self):
+        data = self.dataset.data
+        num_nodes = data.num_nodes
+        num_classes = data.num_classes
+
+        data.add_remaining_self_loops()
+        train_nodes = torch.where(data.train_mask)[0]
+        if len(train_nodes) / num_nodes > self.label_rate:
+            perm = np.random.permutation(train_nodes.shape[0])
+            preserve_nnz = int(num_nodes * self.label_rate)
+            preserved = train_nodes[perm[:preserve_nnz]]
+            masked = train_nodes[perm[preserve_nnz:]]
+            data.train_mask = torch.full((data.train_mask.shape[0],), False, dtype=torch.bool)
+            data.train_mask[preserved] = True
+            data.test_mask[masked] = True
+
+        # Compute absorption probability
+        row, col = data.edge_index
+        A = sp.coo_matrix(
+            (np.ones(row.shape[0]), (row.numpy(), col.numpy())),
+            shape=(num_nodes, num_nodes),
+        ).tocsr()
+        D = A.sum(1).flat
+        confidence = np.zeros([num_classes, num_nodes])
+        confidence_ranking = np.zeros([num_classes, num_nodes], dtype=int)
+
+        if self.approximate:
+            eps = 1e-2
+            for i in range(num_classes):
+                q = list(torch.where(data.y == i)[0].numpy())
+                q = list(filter(lambda x: data.train_mask[x], q))
+                r = {idx: 1 for idx in q}
+                while len(q) > 0:
+                    unode = q.pop()
+                    res = self.alpha / (self.alpha + D[unode]) * r[unode] if unode in r else 0
+                    confidence[i][unode] += res
+                    r[unode] = 0
+                    for vnode in A.indices[A.indptr[unode] : A.indptr[unode + 1]]:
+                        val = res / self.alpha
+                        if vnode in r:
+                            r[vnode] += val
+                        else:
+                            r[vnode] = val
+                        # print(vnode, val)
+                        if val > eps * D[vnode] and vnode not in q:
+                            q.append(vnode)
+        else:
+            L = sp.diags(D, dtype=np.float32) - A
+            L += self.alpha * sp.eye(L.shape[0], dtype=L.dtype)
+            P = slinalg.inv(L.tocsc()).toarray().transpose()
+            for i in range(num_nodes):
+                if data.train_mask[i]:
+                    confidence[data.y[i]] += P[i]
+
+        # Sort nodes by confidence for each class
+        for i in range(num_classes):
+            confidence_ranking[i] = np.argsort(-confidence[i])
+            print(confidence_ranking[i][:10])
+        data.confidence_ranking = confidence_ranking
+
+        self.dataset.data = data
+
+    def pre_stage(self, stage, model_w_out):
+        self.dataset.data.store("y")
+        self.dataset.data.y = model_w_out
+
+    def post_stage(self, stage, model_w_out):
+        self.dataset.data.restore("y")
+
+    def get_dataset(self):
+        return self.dataset
diff --git a/cogdl/wrappers/data_wrapper/node_classification/network_embedding_dw.py b/cogdl/wrappers/data_wrapper/node_classification/network_embedding_dw.py
new file mode 100644
index 00000000..3e1e56ad
--- /dev/null
+++ b/cogdl/wrappers/data_wrapper/node_classification/network_embedding_dw.py
@@ -0,0 +1,26 @@
+import numpy as np
+
+from .. import register_data_wrapper, DataWrapper
+
+
+@register_data_wrapper("network_embedding_dw")
+class NetworkEmbeddingDataWrapper(DataWrapper):
+    def __init__(self, dataset):
+        super(NetworkEmbeddingDataWrapper, self).__init__()
+
+        self.dataset = dataset
+        data = dataset[0]
+
+        num_nodes = data.num_nodes
+        num_classes = dataset.num_classes
+        if len(data.y.shape) > 1:
+            self.label_matrix = data.y
+        else:
+            self.label_matrix = np.zeros((num_nodes, num_classes), dtype=int)
+            self.label_matrix[range(num_nodes), data.y.numpy()] = 1
+
+    def train_wrapper(self):
+        return self.dataset.data
+
+    def test_wrapper(self):
+        return self.label_matrix
diff --git a/cogdl/wrappers/data_wrapper/node_classification/node_classification_dw.py b/cogdl/wrappers/data_wrapper/node_classification/node_classification_dw.py
new file mode 100644
index 00000000..5da2cb8f
--- /dev/null
+++ b/cogdl/wrappers/data_wrapper/node_classification/node_classification_dw.py
@@ -0,0 +1,21 @@
+from .. import DataWrapper, register_data_wrapper
+from cogdl.data import Graph
+
+
+@register_data_wrapper("node_classification_dw")
+class FullBatchNodeClfDataWrapper(DataWrapper):
+    def __init__(self, dataset):
+        super(FullBatchNodeClfDataWrapper, self).__init__(dataset)
+        self.dataset = dataset
+
+    def train_wrapper(self) -> Graph:
+        return self.dataset.data
+
+    def val_wrapper(self):
+        return self.dataset.data
+
+    def test_wrapper(self):
+        return self.dataset.data
+
+    def pre_transform(self):
+        self.dataset.data.add_remaining_self_loops()
diff --git a/cogdl/wrappers/data_wrapper/node_classification/pprgo_dw.py b/cogdl/wrappers/data_wrapper/node_classification/pprgo_dw.py
new file mode 100644
index 00000000..8045d579
--- /dev/null
+++ b/cogdl/wrappers/data_wrapper/node_classification/pprgo_dw.py
@@ -0,0 +1,131 @@
+import os
+import scipy.sparse as sp
+import torch
+
+from .. import DataWrapper, register_data_wrapper
+from cogdl.utils.ppr_utils import build_topk_ppr_matrix_from_data
+
+
+@register_data_wrapper("ppr_dw")
+class PPRDataWrapper(DataWrapper):
+    @staticmethod
+    def add_args(parser):
+        # fmt: off
+        parser.add_argument("--alpha", type=float, default=0.5)
+        parser.add_argument("--topk", type=int, default=32)
+        parser.add_argument("--norm", type=str, default="sym")
+        parser.add_argument("--eps", type=float, default=1e-4)
+
+        parser.add_argument("--batch-size", type=int, default=512)
+        parser.add_argument("--test-batch-size", type=int, default=-1)
+        # fmt: on
+
+    def __init__(self, dataset, topk, alpha=0.2, norm="sym", batch_size=512, eps=1e-4, test_batch_size=-1):
+        super(PPRDataWrapper, self).__init__(dataset)
+        self.batch_size, self.test_batch_size = batch_size, test_batch_size
+        self.topk, self.alpha, self.norm, self.eps = topk, alpha, norm, eps
+        self.dataset = dataset
+
+    def train_wrapper(self):
+        """
+        batch: tuple(x, targets, ppr_scores, y)
+            x: shape=(b, num_features)
+            targets: shape=(num_edges_of_batch,)
+             ppr_scores: shape=(num_edges_of_batch,)
+             y: shape=(b, num_classes)
+        """
+        self.dataset.data.train()
+        ppr_dataset_train = pre_transform(self.dataset, self.topk, self.alpha, self.epsilon, self.norm, mode="train")
+        train_loader = setup_dataloader(ppr_dataset_train, self.batch_size)
+        return train_loader
+
+    def val_wrapper(self):
+        self.dataset.data.eval()
+        if self.test_batch_size > 0:
+            ppr_dataset_val = pre_transform(self.dataset, self.topk, self.alpha, self.epsilon, self.norm, mode="val")
+            val_loader = setup_dataloader(ppr_dataset_val, self.test_batch_size)
+            return val_loader
+        else:
+            return self.dataset.data
+
+    def test_wrapper(self):
+        self.dataset.data.eval()
+        if self.test_batch_size > 0:
+            ppr_dataset_test = pre_transform(self.dataset, self.topk, self.alpha, self.epsilon, self.norm, mode="test")
+            test_loader = setup_dataloader(ppr_dataset_test, self.test_batch_size)
+            return test_loader
+        else:
+            return self.dataset.data
+
+
+def setup_dataloader(ppr_dataset, batch_size):
+    data_loader = torch.utils.data.DataLoader(
+        dataset=ppr_dataset,
+        sampler=torch.utils.data.BatchSampler(
+            torch.utils.data.SequentialSampler(ppr_dataset),
+            batch_size=batch_size,
+            drop_last=False,
+        ),
+        batch_size=None,
+    )
+    return data_loader
+
+
+def pre_transform(dataset, topk, alpha, epsilon, normalization, mode="train"):
+    dataset_name = dataset.__class__.__name__
+    data = dataset[0]
+    num_nodes = data.x.shape[0]
+    nodes = torch.arange(num_nodes)
+
+    mask = getattr(data, f"{mode}_mask")
+    index = nodes[mask].numpy()
+    if mode == "train":
+        data.train()
+    else:
+        data.eval()
+    edge_index = data.edge_index
+
+    if not os.path.exists("./pprgo_saved"):
+        os.mkdir("pprgo_saved")
+    path = f"./pprgo_saved/{dataset_name}_{topk}_{alpha}_{normalization}.{mode}.npz"
+
+    if os.path.exists(path):
+        print(f"Load {mode} from cached")
+        topk_matrix = sp.load_npz(path)
+    else:
+        print(f"Fail to load {mode}, generating...")
+        topk_matrix = build_topk_ppr_matrix_from_data(edge_index, alpha, epsilon, index, topk, normalization)
+        sp.save_npz(path, topk_matrix)
+    result = PPRGoDataset(data.x, topk_matrix, index, data.y)
+    return result
+
+
+class PPRGoDataset(torch.utils.data.Dataset):
+    def __init__(
+        self,
+        features: torch.Tensor,
+        ppr_matrix: sp.csr_matrix,
+        node_indices: torch.Tensor,
+        labels_all: torch.Tensor = None,
+    ):
+        self.features = features
+        self.matrix = ppr_matrix
+        self.node_indices = node_indices
+        self.labels_all = labels_all
+        self.cache = dict()
+
+    def __len__(self):
+        return self.node_indices.shape[0]
+
+    def __getitem__(self, items):
+        key = str(items)
+        if key not in self.cache:
+            sample_matrix = self.matrix[items]
+            source, neighbor = sample_matrix.nonzero()
+            ppr_scores = torch.from_numpy(sample_matrix.data).float()
+
+            features = self.features[neighbor].float()
+            targets = torch.from_numpy(source).long()
+            labels = self.labels_all[self.node_indices[items]]
+            self.cache[key] = (features, targets, ppr_scores, labels)
+        return self.cache[key]
diff --git a/cogdl/wrappers/data_wrapper/node_classification/sagn_dw.py b/cogdl/wrappers/data_wrapper/node_classification/sagn_dw.py
new file mode 100644
index 00000000..35a7caf5
--- /dev/null
+++ b/cogdl/wrappers/data_wrapper/node_classification/sagn_dw.py
@@ -0,0 +1,85 @@
+import torch
+from torch.utils.data import DataLoader
+
+from .. import register_data_wrapper, DataWrapper
+from cogdl.models.nn.sagn import prepare_labels, prepare_feats
+
+
+class SAGEDataWrapper(DataWrapper):
+    @staticmethod
+    def add_args(parser):
+        # fmt: off
+        parser.add_argument("--batch-size", type=int, default=128)
+        parser.add_argument("--label-nhop", type=int, default=3)
+        parser.add_argument("--threshold", type=float, default=0.3)
+        parser.add_argument("--nhtop")
+        # fmt: on
+
+    def __init__(self, dataset, batch_size, label_nhop, threshold, nhtop):
+        super(SAGEDataWrapper, self).__init__(dataset)
+        self.dataset = dataset
+        self.batch_size = batch_size
+        self.label_nhop = label_nhop
+        self.nhop = nhtop
+        self.threshold = threshold
+
+        self.label_emb, self.labels_with_pseudos, self.probs = None, None, None
+        self.multihop_feats = None
+        self.train_nid_with_pseudos = self.dataset.data.train_nid
+
+        self.refresh_per_epoch("train")
+
+    def train_wrapper(self):
+        return DataLoader(self.train_nid_with_pseudos, batch_size=self.batch_size, shuffle=False)
+
+    def val_wrapper(self):
+        val_nid = self.dataset.data.val_nid
+        return DataLoader(val_nid, batch_size=self.batch_size, shuffle=False)
+
+    def test_wrapper(self):
+        test_nid = self.dataset.data
+        return DataLoader(test_nid, batch_size=self.batch_size, shuffle=False)
+
+    def post_stage_wrapper(self):
+        data = self.dataset.data
+        train_nid, val_nid, test_nid = data.train_nid, data.val_nid, data.test_nid
+        all_nid = torch.cat([train_nid, val_nid, test_nid])
+        return DataLoader(all_nid.numpy(), batch_size=self.batch_size, shuffle=False)
+
+    def post_stage_transform(self, batch):
+        return self.train_transform(batch)
+
+    def pre_transform(self):
+        self.multihop_feats = prepare_feats(self.dataset, self.nhop)
+
+    def train_transform(self, batch):
+        batch_x = [x[batch] for x in self.multihop_feats]
+        if self.label_emb is not None:
+            batch_y_emb = self.label_emb[batch]
+        else:
+            batch_y_emb = None
+        y = self.labels_with_pseudos[batch]
+        return [batch_x, batch_y_emb, y]
+
+    def val_transform(self, batch):
+        batch_x = [x[batch] for x in self.multihop_feats]
+        if self.label_emb is not None:
+            batch_y_emb = self.label_emb[batch]
+        else:
+            batch_y_emb = None
+        y = self.dataset.data.y[batch]
+        return [batch_x, batch_y_emb, y]
+
+    def test_transform(self, batch):
+        return self.val_transform(batch)
+
+    def pre_stage(self, stage, model_w_out):
+        dataset = self.dataset
+        probs = model_w_out
+        with torch.no_grad():
+            (label_emb, labels_with_pseudos, train_nid_with_pseudos) = prepare_labels(
+                dataset, stage, self.label_nhop, self.threshold, probs=probs
+            )
+        self.label_emb = label_emb
+        self.labels_with_pseudos = labels_with_pseudos
+        self.train_nid_with_pseudos = train_nid_with_pseudos
diff --git a/cogdl/wrappers/data_wrapper/pretraining/__init__.py b/cogdl/wrappers/data_wrapper/pretraining/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/cogdl/wrappers/data_wrapper/pretraining/gcc_dw.py b/cogdl/wrappers/data_wrapper/pretraining/gcc_dw.py
new file mode 100644
index 00000000..92029211
--- /dev/null
+++ b/cogdl/wrappers/data_wrapper/pretraining/gcc_dw.py
@@ -0,0 +1,289 @@
+import copy
+import math
+from typing import Tuple
+
+from scipy.sparse import linalg
+
+import numpy as np
+import scipy.sparse as sparse
+import sklearn.preprocessing as preprocessing
+import torch
+import torch.nn.functional as F
+
+from torch.utils.data import DataLoader
+
+from .. import register_data_wrapper, DataWrapper
+from cogdl.data import batch_graphs, Graph
+
+
+@register_data_wrapper("gcc_dw")
+class GCCDataWrapper(DataWrapper):
+    @staticmethod
+    def add_args(parser):
+        # random walk
+        parser.add_argument("--batch-size", type=int, default=128)
+        parser.add_argument("--rw-hops", type=int, default=64)
+        parser.add_argument("--subgraph-size", type=int, default=128)
+        parser.add_argument("--restart-prob", type=float, default=0.8)
+        parser.add_argument("--positional-embedding-size", type=int, default=128)
+        parser.add_argument(
+            "--task", type=str, default="node_classification", choices=["node_classification, graph_classification"]
+        )
+        parser.add_argument("--num-workers", type=int, default=4)
+
+    def __init__(
+        self,
+        dataset,
+        batch_size,
+        finetune=False,
+        num_workers=4,
+        rw_hops=64,
+        subgraph_size=128,
+        restart_prob=0.8,
+        positional_embedding_size=128,
+        task="node_classification",
+    ):
+        super(GCCDataWrapper, self).__init__(dataset)
+
+        data = dataset.data
+        data.add_remaining_self_loops()
+        if task == "node_classification":
+            if finetune:
+                self.train_dataset = NodeClassificationDatasetLabeled(
+                    data, rw_hops, subgraph_size, restart_prob, positional_embedding_size
+                )
+            else:
+                self.train_dataset = NodeClassificationDataset(
+                    data, rw_hops, subgraph_size, restart_prob, positional_embedding_size
+                )
+        elif task == "graph_classification":
+            if finetune:
+                pass
+            else:
+                pass
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        self.finetune = finetune
+
+        self.rw_hops = rw_hops
+        self.subgraph_size = subgraph_size
+        self.restart_prob = restart_prob
+
+    def train_wrapper(self):
+        train_loader = DataLoader(
+            dataset=self.train_dataset,
+            batch_size=self.batch_size,
+            collate_fn=labeled_batcher() if self.finetune else batcher(),
+            shuffle=True if self.finetune else False,
+            num_workers=self.num_workers,
+            worker_init_fn=None,
+        )
+        return train_loader
+
+
+def labeled_batcher():
+    def batcher_dev(batch):
+        graph_q, label = zip(*batch)
+        graph_q = batch_graphs(graph_q)
+        return graph_q, torch.LongTensor(label)
+
+    return batcher_dev
+
+
+def batcher():
+    def batcher_dev(batch):
+        graph_q_, graph_k_ = zip(*batch)
+        graph_q, graph_k = batch_graphs(graph_q_), batch_graphs(graph_k_)
+        graph_q.batch_size = len(graph_q_)
+        return graph_q, graph_k
+
+    return batcher_dev
+
+
+def eigen_decomposision(n, k, laplacian, hidden_size, retry):
+    if k <= 0:
+        return torch.zeros(n, hidden_size)
+    laplacian = laplacian.astype("float64")
+    ncv = min(n, max(2 * k + 1, 20))
+    # follows https://stackoverflow.com/questions/52386942/scipy-sparse-linalg-eigsh-with-fixed-seed
+    v0 = np.random.rand(n).astype("float64")
+    for i in range(retry):
+        try:
+            s, u = linalg.eigsh(laplacian, k=k, which="LA", ncv=ncv, v0=v0)
+        except sparse.linalg.eigen.arpack.ArpackError:
+            # print("arpack error, retry=", i)
+            ncv = min(ncv * 2, n)
+            if i + 1 == retry:
+                sparse.save_npz("arpack_error_sparse_matrix.npz", laplacian)
+                u = torch.zeros(n, k)
+        else:
+            break
+    x = preprocessing.normalize(u, norm="l2")
+    x = torch.from_numpy(x.astype("float32"))
+    x = F.pad(x, (0, hidden_size - k), "constant", 0)
+    return x
+
+
+def _add_undirected_graph_positional_embedding(g: Graph, hidden_size, retry=10):
+    # We use eigenvectors of normalized graph laplacian as vertex features.
+    # It could be viewed as a generalization of positional embedding in the
+    # attention is all you need paper.
+    # Recall that the eignvectors of normalized laplacian of a line graph are cos/sin functions.
+    # See section 2.4 of http://www.cs.yale.edu/homes/spielman/561/2009/lect02-09.pdf
+    n = g.num_nodes
+    with g.local_graph():
+        g.sym_norm()
+        adj = g.to_scipy_csr()
+    laplacian = adj
+
+    # adj = g.adjacency_matrix_scipy(transpose=False, return_edge_ids=False).astype(float)
+    # norm = sparse.diags(dgl.backend.asnumpy(g.in_degrees()).clip(1) ** -0.5, dtype=float)
+    # laplacian = norm * adj * norm
+
+    k = min(n - 2, hidden_size)
+    x = eigen_decomposision(n, k, laplacian, hidden_size, retry)
+    # g.ndata["pos_undirected"] = x.float()
+    g.pos_undirected = x.float()
+    return g
+
+
+def _rwr_trace_to_cogdl_graph(
+    g: Graph, seed: int, trace: torch.Tensor, positional_embedding_size: int, entire_graph: bool = False
+):
+    subv = torch.unique(trace).tolist()
+    try:
+        subv.remove(seed)
+    except ValueError:
+        pass
+    subv = [seed] + subv
+    if entire_graph:
+        subg = copy.deepcopy(g)
+    else:
+        subg = g.subgraph(subv)
+
+    subg = _add_undirected_graph_positional_embedding(subg, positional_embedding_size)
+
+    subg.seed = torch.zeros(subg.num_nodes, dtype=torch.long)
+    if entire_graph:
+        subg.seed[seed] = 1
+    else:
+        subg.seed[0] = 1
+    return subg
+
+
+class NodeClassificationDataset(object):
+    def __init__(
+        self,
+        data: Graph,
+        rw_hops: int = 64,
+        subgraph_size: int = 64,
+        restart_prob: float = 0.8,
+        positional_embedding_size: int = 32,
+        step_dist: list = [1.0, 0.0, 0.0],
+    ):
+        self.rw_hops = rw_hops
+        self.subgraph_size = subgraph_size
+        self.restart_prob = restart_prob
+        self.positional_embedding_size = positional_embedding_size
+        self.step_dist = step_dist
+        assert positional_embedding_size > 1
+
+        self.data = data
+        self.graphs = [self.data]
+        self.length = sum([g.num_nodes for g in self.graphs])
+        self.total = self.length
+
+    def __len__(self):
+        return self.length
+
+    def _convert_idx(self, idx) -> Tuple[int, int]:
+        graph_idx = 0
+        node_idx = idx
+        for i in range(len(self.graphs)):
+            if node_idx < self.graphs[i].num_nodes:
+                graph_idx = i
+                break
+            else:
+                node_idx -= self.graphs[i].num_nodes
+        return graph_idx, node_idx
+
+    def __getitem__(self, idx):
+        graph_idx, node_idx = self._convert_idx(idx)
+
+        step = np.random.choice(len(self.step_dist), 1, p=self.step_dist)[0]
+        g = self.graphs[graph_idx]
+
+        if step == 0:
+            other_node_idx = node_idx
+        else:
+            other_node_idx = g.random_walk([node_idx], step)[-1]
+
+        max_nodes_per_seed = max(
+            self.rw_hops,
+            int((self.graphs[graph_idx].degrees()[node_idx] * math.e / (math.e - 1) / self.restart_prob) + 0.5),
+        )
+        # TODO: `num_workers > 0` is not compatible with `numba`
+        traces = g.random_walk_with_restart([node_idx, other_node_idx], max_nodes_per_seed, self.restart_prob)
+
+        # traces = [[0,1,2,3], [1,2,3,4]]
+
+        graph_q = _rwr_trace_to_cogdl_graph(
+            g=g,
+            seed=node_idx,
+            trace=torch.Tensor(traces[0]),
+            positional_embedding_size=self.positional_embedding_size,
+            entire_graph=hasattr(self, "entire_graph") and self.entire_graph,
+        )
+        graph_k = _rwr_trace_to_cogdl_graph(
+            g=g,
+            seed=other_node_idx,
+            trace=torch.Tensor(traces[1]),
+            positional_embedding_size=self.positional_embedding_size,
+            entire_graph=hasattr(self, "entire_graph") and self.entire_graph,
+        )
+        return graph_q, graph_k
+
+
+class NodeClassificationDatasetLabeled(NodeClassificationDataset):
+    def __init__(
+        self,
+        data,
+        rw_hops=64,
+        subgraph_size=64,
+        restart_prob=0.8,
+        positional_embedding_size=32,
+        step_dist=[1.0, 0.0, 0.0],
+    ):
+        super(NodeClassificationDatasetLabeled, self).__init__(
+            data,
+            rw_hops,
+            subgraph_size,
+            restart_prob,
+            positional_embedding_size,
+            step_dist,
+        )
+        assert len(self.graphs) == 1
+        self.num_classes = self.data.num_classes
+
+    def __getitem__(self, idx):
+        graph_idx = 0
+        node_idx = idx
+        for i in range(len(self.graphs)):
+            if node_idx < self.graphs[i].num_nodes:
+                graph_idx = i
+                break
+            else:
+                node_idx -= self.graphs[i].num_nodes
+
+        g = self.graphs[graph_idx]
+        traces = g.random_walk_with_restart([node_idx], self.rw_hops, self.restart_prob)
+
+        graph_q = _rwr_trace_to_cogdl_graph(
+            g=g,
+            seed=node_idx,
+            trace=torch.Tensor(traces[0]),
+            positional_embedding_size=self.positional_embedding_size,
+        )
+        graph_q.y = self.data.y[idx].y
+        return graph_q
+        # return graph_q, self.data.y[idx].argmax().item()
diff --git a/cogdl/wrappers/default_match.py b/cogdl/wrappers/default_match.py
new file mode 100644
index 00000000..0065c07b
--- /dev/null
+++ b/cogdl/wrappers/default_match.py
@@ -0,0 +1,122 @@
+from cogdl.wrappers import fetch_model_wrapper, fetch_data_wrapper
+
+
+def set_default_wrapper_config():
+    node_classification_models = [
+        "gcn",
+        "deepergcn",
+        "drgcn",
+        "drgat",
+        "gcnii",
+        "gcnmix",
+        "grand",
+        "grace",
+        "mvgrl",
+        "graphsage",
+        "gdc_gcn",
+        "mixhop",
+        "mlp",
+        "moe_gcn",
+        "ppnp",
+        "appnp",
+        "pprgo",
+        "chebyshev",
+        "pyg_gcn",
+        "unet",
+        "srgcn",
+        "revgcn",
+        "revgat",
+        "revgen",
+        "sagn",
+        "sign",
+        "sgc",
+        "unsup_graphsage",
+        "dgi",
+        "dropedge_gcn",
+        "gat",
+        "graphsaint",
+        "m3s",
+        "pyg_supergat",
+    ]
+
+    graph_classification_models = [
+        "gin",
+        "pathcy_san",
+        "diffpool",
+        "infograph",
+        "sagpool",
+    ]
+
+    network_embedding_models = [
+        "deepwalk",
+        "line",
+        "node2veec",
+        "prone",
+        "netmf",
+        "netsmf",
+        "sdne",
+        "spectral",
+        "dngr",
+        "grarep",
+        "hope",
+    ]
+
+    graph_embedding_models = [
+        "dgk",
+        "graph2vec",
+    ]
+
+    node_classification_wrappers = dict()
+    for item in node_classification_models:
+        node_classification_wrappers[item] = {"mw": "node_classification_mw", "dw": "node_classification_dw"}
+
+    node_classification_wrappers["dgi"]["mw"] = "dgi_mw"
+    node_classification_wrappers["m3s"]["mw"] = "m3s_mw"
+    node_classification_wrappers["graphsage"]["mw"] = "graphsage_mw"
+    node_classification_wrappers["mvgrl"]["mw"] = "mvgrl_mw"
+    node_classification_wrappers["sagn"]["mw"] = "sagn_mw"
+    node_classification_wrappers["grand"]["mw"] = "grand_mw"
+    node_classification_wrappers["gcnmix"]["mw"] = "gcnmix_mw"
+
+    node_classification_wrappers["m3s"]["dw"] = "m3s_dw"
+    node_classification_wrappers["graphsage"]["dw"] = "graphsage_dw"
+    node_classification_wrappers["pprgo"]["dw"] = "pprgo_dw"
+    node_classification_wrappers["sagn"]["dw"] = "sagn_dw"
+
+    graph_classification_wrappers = dict()
+    for item in graph_classification_models:
+        graph_classification_wrappers[item] = {"mw": "graph_classification_mw", "dw": "graph_classification_dw"}
+
+    graph_classification_wrappers["infograph"]["mw"] = "infograph_mw"
+
+    network_embedding_wrappers = dict()
+    for item in network_embedding_models:
+        network_embedding_wrappers[item] = {"mw": "network_embedding_mw", "dw": "network_embedding_dw"}
+
+    graph_embedding_wrappers = dict()
+    for item in graph_embedding_models:
+        graph_embedding_wrappers[item] = {"mw": "graph_embedding_mw", "dw": "graph_embedding_dw"}
+
+    merged = dict()
+    merged.update(node_classification_wrappers)
+    merged.update(graph_embedding_wrappers)
+    merged.update(graph_embedding_wrappers)
+    merged.update(network_embedding_wrappers)
+    return merged
+
+
+default_wrapper_config = set_default_wrapper_config()
+
+
+def get_wrappers(model_name):
+    if model_name in default_wrapper_config:
+        dw = default_wrapper_config[model_name]["dw"]
+        mw = default_wrapper_config[model_name]["mw"]
+        return fetch_model_wrapper(mw), fetch_data_wrapper(dw)
+
+
+def get_wrappers_name(model_name):
+    if model_name in default_wrapper_config:
+        dw = default_wrapper_config[model_name]["dw"]
+        mw = default_wrapper_config[model_name]["mw"]
+        return mw, dw
diff --git a/cogdl/wrappers/model_wrapper/__init__.py b/cogdl/wrappers/model_wrapper/__init__.py
new file mode 100644
index 00000000..24de739a
--- /dev/null
+++ b/cogdl/wrappers/model_wrapper/__init__.py
@@ -0,0 +1,63 @@
+import os
+import importlib
+
+from .base_model_wrapper import ModelWrapper, EmbeddingModelWrapper
+
+
+MODELMODULE_REGISTRY = {}
+SUPPORTED_MODELMODULES = {}
+
+
+def register_model_wrapper(name):
+    """
+    New data wrapper types can be added to cogdl with the :func:`register_model_wrapper`
+    function decorator.
+
+    Args:
+        name (str): the name of the model_wrapper
+    """
+
+    def register_model_wrapper_cls(cls):
+        if name in MODELMODULE_REGISTRY:
+            raise ValueError("Cannot register duplicate model_wrapper ({})".format(name))
+        if not issubclass(cls, ModelWrapper):
+            raise ValueError("Model ({}: {}) must extend BaseModel".format(name, cls.__name__))
+        MODELMODULE_REGISTRY[name] = cls
+        cls.model_name = name
+        return cls
+
+    return register_model_wrapper_cls
+
+
+def scan_model_wrappers():
+    global SUPPORTED_MODELMODULES
+    dirname = os.path.dirname(__file__)
+    dir_names = [x for x in os.listdir(dirname) if not x.startswith("__")]
+    dirs = [os.path.join(dirname, x) for x in dir_names]
+    dirs_names = [(x, y) for x, y in zip(dirs, dir_names) if os.path.isdir(x)]
+    dw_dict = SUPPORTED_MODELMODULES
+    for _dir, _name in dirs_names:
+        files = os.listdir(_dir)
+        # files = [x for x in os.listdir(_dir) if os.path.isfile(x)]
+        dw = [x.split(".")[0] for x in files]
+        dw = [x for x in dw if not x.startswith("__")]
+        path = [f"cogdl.wrappers.model_wrapper.{_name}.{x}" for x in dw]
+        for x, y in zip(dw, path):
+            dw_dict[x] = y
+
+
+def try_import_model_wrapper(name):
+    if name in MODELMODULE_REGISTRY:
+        return
+    if name in SUPPORTED_MODELMODULES:
+        importlib.import_module(SUPPORTED_MODELMODULES[name])
+    else:
+        raise NotImplementedError(f"`{name}` model_wrapper is not implemented.")
+
+
+def fetch_model_wrapper(name):
+    try_import_model_wrapper(name)
+    return MODELMODULE_REGISTRY[name]
+
+
+scan_model_wrappers()
diff --git a/cogdl/wrappers/model_wrapper/base_model_wrapper.py b/cogdl/wrappers/model_wrapper/base_model_wrapper.py
new file mode 100644
index 00000000..1b89616c
--- /dev/null
+++ b/cogdl/wrappers/model_wrapper/base_model_wrapper.py
@@ -0,0 +1,201 @@
+from typing import Union, Callable
+from abc import abstractmethod
+import torch
+from cogdl.wrappers.tools.wrapper_utils import merge_batch_indexes
+from cogdl.utils.evaluator import setup_evaluator, Accuracy, MultiLabelMicroF1
+
+
+class ModelWrapper(torch.nn.Module):
+    @staticmethod
+    def add_args(parser):
+        pass
+
+    def __init__(self):
+        super(ModelWrapper, self).__init__()
+        self.__model_keys__ = None
+        self._loss_func = None
+        self._evaluator = None
+        self._evaluator_metric = None
+        self.__record__ = dict()
+
+    def forward(self):
+        pass
+
+    def pre_stage(self, stage, data_w):
+        pass
+
+    def post_stage(self, stage, data_w):
+        pass
+
+    def train_step(self, subgraph):
+        pass
+
+    def val_step(self, subgraph):
+        pass
+
+    def test_step(self, subgraph):
+        pass
+
+    def evaluate(self, pred: torch.Tensor, labels: torch.Tensor, metric: Union[str, Callable] = "auto"):
+        """
+        method: str or callable function,
+        """
+        pred = pred.cpu()
+        labels = labels.cpu()
+        if self._evaluator is None:
+            if metric == "auto":
+                if len(labels.shape) > 1:
+                    metric = "multilabel_microf1"
+                    self._evaluator_metric = "micro_f1"
+                else:
+                    metric = "accuracy"
+                    self._evaluator_metric = "acc"
+
+            self._evaluator = setup_evaluator(metric)
+            # self._evaluator_metric = metric
+        return self._evaluator(pred, labels)
+
+    @abstractmethod
+    def setup_optimizer(self):
+        raise NotImplementedError
+
+    def set_early_stopping(self):
+        """
+        Return:
+            1. `str`, the monitoring metric
+            2. tuple(`str`, `str`), that is, (the monitoring metric, `small` or `big`). The second parameter means,
+                `the smaller, the better` or `the bigger, the better`
+        """
+        return "val_metric", ">"
+
+    def on_train_step(self, *args, **kwargs):
+        return self.train_step(*args, **kwargs)
+
+    def on_val_step(self, *args, **kwargs):
+        out = self.val_step(*args, **kwargs)
+        self.set_notes(out, "val")
+
+    def on_test_step(self, *args, **kwargs):
+        out = self.test_step(*args, **kwargs)
+        self.set_notes(out, "test")
+
+    # def on_evaluate(self, *args, **kwargs):
+    #     out = self.evaluate(*args, **kwargs)
+    #     self.set_notes(out, "evaluate")
+
+    def set_notes(self, out, prefix="val"):
+        if isinstance(out, dict):
+            for key, val in out.items():
+                self.note(key, val)
+        elif isinstance(out, tuple) or isinstance(out, list):
+            for i, val in enumerate(out):
+                self.note(f"{prefix}_{i}", val)
+
+    def note(self, name: str, data, merge="mean"):
+        """
+        name: str
+        data: Any
+        """
+        if name not in self.__record__:
+            name = name.lower()
+            self.__record__[name] = [data]
+            # self.__record_merge__[name] = merge
+        else:
+            self.__record__[name].append(data)
+
+    def collect_notes(self):
+        if len(self.__record__) == 0:
+            return None
+        out = dict()
+        for key, val in self.__record__.items():
+            if key.endswith("_metric"):
+                _val = self._evaluator.evaluate()
+
+            elif isinstance(self._evaluator_metric, str) and key.endswith(self._evaluator_metric):
+                _val = self._evaluator.evaluate()
+            else:
+                _val = merge_batch_indexes(val)
+            out[key] = _val
+        self.__record__ = dict()
+        return out
+
+    @property
+    def default_loss_fn(self):
+        if self._loss_func is None:
+            raise RuntimeError(
+                "`loss_fn` must be set for your ModelWrapper using `mw.default_loss_fn = your_loss_fn`.",
+                f"Now self.loss_fn is {self._loss_fn}",
+            )
+        return self._loss_func
+
+    @default_loss_fn.setter
+    def default_loss_fn(self, loss_fn):
+        self._loss_func = loss_fn
+
+    @property
+    def default_evaluator(self):
+        return self._evaluator
+
+    @default_evaluator.setter
+    def default_evaluator(self, x):
+        self._evaluator = x
+
+    @property
+    def device(self):
+        # for k in self._model_key_:
+        #     return next(getattr(self, k).parameters()).device
+        return next(self.parameters()).device
+
+    @property
+    def evaluation_metric(self):
+        return self._evaluator_metric
+
+    def set_evaluation_metric(self):
+        if isinstance(self._evaluator, MultiLabelMicroF1):
+            self._evaluator_metric = "micro_f1"
+        elif isinstance(self._evaluator, Accuracy):
+            self._evaluator_metric = "acc"
+        else:
+            evaluation_metric = self.set_early_stopping()
+            if not isinstance(evaluation_metric, str):
+                evaluation_metric = evaluation_metric[0]
+            if evaluation_metric.startswith("val"):
+                evaluation_metric = evaluation_metric[3:]
+            self._evaluator_metric = evaluation_metric
+
+    def load_checkpoint(self, path):
+        pass
+
+    def save_checkpoint(self, path):
+        pass
+
+    def _find_model(self):
+        models = []
+        for k, v in self.__dict__.items():
+            if isinstance(v, torch.nn.Module):
+                models.append(k)
+        self.__model_keys__ = models
+
+    @property
+    def wrapped_model(self):
+        if hasattr(self, "model"):
+            return getattr(self, "model")
+        assert len(self._model_key_) == 1, f"{len(self._model_key_)} exists"
+        return getattr(self, self._model_key_[0])
+
+    @wrapped_model.setter
+    def wrapped_model(self, model):
+        if len(self._model_key_) == 0:
+            self.__model_keys__ = [None]
+        setattr(self, self._model_key_[0], model)
+
+    @property
+    def _model_key_(self):
+        if self.__model_keys__ is None:
+            self._find_model()
+        return self.__model_keys__
+
+
+class EmbeddingModelWrapper(ModelWrapper):
+    def setup_optimizer(self):
+        pass
diff --git a/cogdl/wrappers/model_wrapper/clustering/__init__.py b/cogdl/wrappers/model_wrapper/clustering/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/cogdl/wrappers/model_wrapper/clustering/daegc_mw.py b/cogdl/wrappers/model_wrapper/clustering/daegc_mw.py
new file mode 100644
index 00000000..14148eb0
--- /dev/null
+++ b/cogdl/wrappers/model_wrapper/clustering/daegc_mw.py
@@ -0,0 +1,109 @@
+import torch
+import torch.nn.functional as F
+
+from sklearn.cluster import KMeans
+
+from .. import register_model_wrapper, ModelWrapper
+from cogdl.wrappers.tools.wrapper_utils import evaluate_clustering
+
+
+@register_model_wrapper("dagae_mw")
+class DAGAEModelWrapper(ModelWrapper):
+    @staticmethod
+    def add_args(parser):
+        # fmt: off
+        parser.add_argument("--num-clusters", type=int, default=7)
+        parser.add_argument("--cluster-method", type=str, default="kmeans", help="option: kmeans or spectral")
+        parser.add_argument("--evaluation", type=str, default="full", help="option: full or NMI")
+        parser.add_argument("--T", type=int, default=5)
+        # fmt: on
+
+    def __init__(self, model, optimizer_cfg, num_clusters, cluster_method="kmeans", evaluation="full", T=5):
+        super(DAGAEModelWrapper, self).__init__()
+        self.model = model
+        self.num_clusters = num_clusters
+        self.optimizer_cfg = optimizer_cfg
+        self.cluster_method = cluster_method
+        self.full = evaluation == "full"
+        self.t = T
+
+        self.stage = 0
+        self.count = 0
+
+    def train_step(self, subgraph):
+        graph = subgraph
+        if self.stage == 0:
+            z = self.model(graph)
+            loss = self.recon_loss(z, graph.adj_mx)
+        else:
+            cluster_center = self.model.get_cluster_center()
+            z = self.model(graph)
+            Q = self.getQ(z, cluster_center)
+            # if epoch % self.T == 0:
+            self.count += 1
+            if self.count % self.t == 0:
+                P = self.getP(Q).detach()
+            loss = self.recon_loss(z, graph.adj_mx) + self.gamma * self.cluster_loss(P, Q)
+        return loss
+
+    def evaluate(self, dataset):
+        data = dataset.data
+        features_matrix = self.model(data)
+        features_matrix = features_matrix.cpu().numpy()
+        return evaluate_clustering(
+            features_matrix, data.y, self.cluster_method, self.num_clusters, data.num_nodes, self.full
+        )
+
+    def recon_loss(self, z, adj):
+        return F.binary_cross_entropy(F.softmax(torch.mm(z, z.t())), adj, reduction="sum")
+
+    def cluster_loss(self, P, Q):
+        # return nn.MSELoss(reduce=True, size_average=False)(P, Q)
+        return torch.nn.KLDivLoss(reduce=True, size_average=False)(P.log(), Q)
+
+    def setup_optimizer(self):
+        lr, wd = self.optimizer_cfg["lr"], self.optimizer_cfg["weight_decay"]
+        return torch.optim.Adam(self.parameters(), lr=lr, weight_decay=wd)
+
+    def pre_stage(self, stage, data_w):
+        self.stage = stage
+        if stage == 0:
+            data = data_w.get_dataset().data
+            data.add_remaining_self_loops()
+
+            data.store("edge_index")
+
+            data.adj_mx = torch.sparse_coo_tensor(
+                torch.stack(data.edge_index),
+                torch.ones(data.edge_index[0].shape[0]),
+                torch.Size([data.x.shape[0], data.x.shape[0]]),
+            ).to_dense()
+            edge_index_2hop = data.edge_index
+            data.edge_index = edge_index_2hop
+
+    def post_stage(self, stage, data_w):
+        if stage == 0:
+            data = data_w.get_dataset().data
+            data.restore("edge_index")
+            data.to(self.device)
+            kmeans = KMeans(n_clusters=self.num_clusters, random_state=0).fit(self.model(data).detach().cpu().numpy())
+            self.model.set_cluster_center(torch.tensor(kmeans.cluster_centers_, device=self.device))
+
+    def getQ(self, z, cluster_center):
+        Q = None
+        for i in range(z.shape[0]):
+            dis = torch.sum((z[i].repeat(self.num_clusters, 1) - 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)
+        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
diff --git a/cogdl/wrappers/model_wrapper/clustering/gae_mw.py b/cogdl/wrappers/model_wrapper/clustering/gae_mw.py
new file mode 100644
index 00000000..161502f2
--- /dev/null
+++ b/cogdl/wrappers/model_wrapper/clustering/gae_mw.py
@@ -0,0 +1,59 @@
+import torch
+import torch.nn.functional as F
+
+from .. import register_model_wrapper, ModelWrapper
+from cogdl.wrappers.tools.wrapper_utils import evaluate_clustering
+
+
+@register_model_wrapper("gat_mw")
+class GAEModelWrapper(ModelWrapper):
+    @staticmethod
+    def add_args(parser):
+        # fmt: off
+        parser.add_argument("--num-clusters", type=int, default=7)
+        parser.add_argument("--cluster-method", type=str, default="kmeans", help="option: kmeans or spectral")
+        parser.add_argument("--evaluation", type=str, default="full", help="option: full or NMI")
+        # fmt: on
+
+    def __init__(self, model, optimizer_cfg, num_clusters, cluster_method="kmeans", evaluation="full"):
+        super(GAEModelWrapper, self).__init__()
+        self.model = model
+        self.num_clusters = num_clusters
+        self.optimizer_cfg = optimizer_cfg
+        self.cluster_method = cluster_method
+        self.full = evaluation == "full"
+
+    def train_step(self, subgraph):
+        graph = subgraph
+        mean, log_var = self.model.encode(graph)
+        z = self.reparameterization(mean, log_var)
+        mat = self.model.decode(z)
+        recon_loss = F.binary_cross_entropy(mat, graph.adj_mx, reduction="sum")
+        var = torch.exp(log_var)
+        kl_loss = 0.5 * torch.mean(torch.sum(mean * mean + var - log_var - 1, dim=1))
+        # print("recon_loss = %.3f, kl_loss = %.3f" % (recon_loss, kl_loss))
+        return recon_loss + kl_loss
+
+    def evaluate(self, dataset):
+        data = dataset.data
+        features_matrix = self.model(data)
+        features_matrix = features_matrix.cpu().numpy()
+        return evaluate_clustering(
+            features_matrix, data.y, self.cluster_method, self.num_clusters, data.num_nodes, self.full
+        )
+
+    @staticmethod
+    def reparameterization(mean, log_var):
+        sigma = torch.exp(log_var)
+        z = mean + torch.randn_like(log_var) * sigma
+        return z
+
+    def pre_stage(self, stage, data_w):
+        if stage == 0:
+            data = data_w.get_dataset().data
+            adj_mx = torch.sparse_coo_tensor(
+                torch.stack(data.edge_index),
+                torch.ones(data.edge_index[0].shape[0]),
+                torch.Size([data.x.shape[0], data.x.shape[0]]),
+            ).to_dense()
+            data.adj_mx = adj_mx
diff --git a/cogdl/wrappers/model_wrapper/graph_classification/__init__.py b/cogdl/wrappers/model_wrapper/graph_classification/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/cogdl/wrappers/model_wrapper/graph_classification/graph_classification_mw.py b/cogdl/wrappers/model_wrapper/graph_classification/graph_classification_mw.py
new file mode 100644
index 00000000..432a2938
--- /dev/null
+++ b/cogdl/wrappers/model_wrapper/graph_classification/graph_classification_mw.py
@@ -0,0 +1,42 @@
+import torch
+
+from .. import register_model_wrapper, ModelWrapper
+from cogdl.wrappers.tools.wrapper_utils import pre_evaluation_index
+
+
+@register_model_wrapper("graph_classification_mw")
+class GraphClassificationModelWrapper(ModelWrapper):
+    def __init__(self, model, optimizer_cfg):
+        super(GraphClassificationModelWrapper, self).__init__()
+        self.model = model
+        self.optimizer_cfg = optimizer_cfg
+
+    def train_step(self, batch):
+        pred = self.model(batch)
+        y = batch.y
+        loss = self.default_loss_fn(pred, y)
+        return loss
+
+    def val_step(self, batch):
+        pred = self.model(batch)
+        y = batch.y
+        val_loss = self.default_loss_fn(pred, y)
+
+        metric = self.evaluate(pred, y, metric="auto")
+
+        self.note("val_loss", val_loss)
+        self.note("val_metric", metric)
+
+    def test_step(self, batch):
+        pred = self.model(batch)
+        y = batch.y
+        test_loss = self.default_loss_fn(pred, y)
+
+        metric = self.evaluate(pred, y, metric="auto")
+
+        self.note("test_loss", test_loss)
+        self.note("test_metric", metric)
+
+    def setup_optimizer(self):
+        cfg = self.optimizer_cfg
+        return torch.optim.Adam(self.model.parameters(), lr=cfg["lr"], weight_decay=cfg["weight_decay"])
diff --git a/cogdl/wrappers/model_wrapper/graph_classification/graph_embedding_mw.py b/cogdl/wrappers/model_wrapper/graph_classification/graph_embedding_mw.py
new file mode 100644
index 00000000..acffa640
--- /dev/null
+++ b/cogdl/wrappers/model_wrapper/graph_classification/graph_embedding_mw.py
@@ -0,0 +1,24 @@
+from torch.utils.data import DataLoader
+
+from cogdl.data import MultiGraphDataset
+from .. import register_model_wrapper, EmbeddingModelWrapper
+from cogdl.wrappers.tools.wrapper_utils import evaluate_graph_embeddings_using_svm
+
+
+@register_model_wrapper("graph_embedding_mw")
+class GraphEmbeddingModelWrapper(EmbeddingModelWrapper):
+    def __init__(self, model):
+        super(GraphEmbeddingModelWrapper, self).__init__()
+        self.model = model
+
+    def train_step(self, batch):
+        if isinstance(batch, DataLoader) or isinstance(batch, MultiGraphDataset):
+            graphs = [x for x in batch]
+        else:
+            graphs = batch
+        emb = self.model(graphs)
+        return emb
+
+    def test_step(self, batch):
+        x, y = batch
+        return evaluate_graph_embeddings_using_svm(x, y)
diff --git a/cogdl/wrappers/model_wrapper/graph_classification/infograph_mw.py b/cogdl/wrappers/model_wrapper/graph_classification/infograph_mw.py
new file mode 100644
index 00000000..f7a660e0
--- /dev/null
+++ b/cogdl/wrappers/model_wrapper/graph_classification/infograph_mw.py
@@ -0,0 +1,125 @@
+import math
+import numpy as np
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .. import register_model_wrapper, ModelWrapper
+from cogdl.models.nn.mlp import MLP
+from cogdl.data import DataLoader
+from cogdl.wrappers.tools.wrapper_utils import evaluate_graph_embeddings_using_svm
+
+
+@register_model_wrapper("infograph_mw")
+class InfoGraphModelWrapper(ModelWrapper):
+    @staticmethod
+    def add_args(parser):
+        # fmt: off
+        parser.add_argument("--sup", action="store_true")
+        # fmt: on
+
+    def __init__(self, model, optimizer_cfg, sup=False):
+        super(InfoGraphModelWrapper, self).__init__()
+        self.model = model
+        hidden_size = optimizer_cfg["hidden_size"]
+        model_num_layers = model.num_layers
+        self.local_dis = FF(model_num_layers * hidden_size, hidden_size)
+        self.global_dis = FF(model_num_layers * hidden_size, hidden_size)
+
+        self.optimizer_cfg = optimizer_cfg
+        self.sup = sup
+        self.criterion = torch.nn.MSELoss()
+
+    def train_step(self, batch):
+        if self.sup:
+            pred = self.model.sup_forward(batch, batch.x)
+            loss = self.sup_loss(pred, batch)
+        else:
+            graph_feat, node_feat = self.model.unsup_forward(batch, batch.x)
+            loss = self.unsup_loss(graph_feat, node_feat, batch.batch)
+        return loss
+
+    def test_step(self, dataset):
+        device = self.device
+        dataloader = DataLoader(dataset, batch_size=32, shuffle=False)
+        preds = []
+        with torch.no_grad():
+            for batch in dataloader:
+                preds.append(self.model(batch.to(device)))
+        preds = torch.cat(preds).cpu().numpy()
+        labels = np.array([g.y.item() for g in dataset])
+        result = evaluate_graph_embeddings_using_svm(preds, labels)
+
+        self.note("test_acc", result["acc"])
+        self.note("std", result["std"])
+
+    def setup_optimizer(self):
+        cfg = self.optimizer_cfg
+        return torch.optim.Adam(
+            [
+                {"params": self.model.parameters()},
+                {"params": self.global_dis.parameters()},
+                {"params": self.local_dis.parameters()},
+            ],
+            lr=cfg["lr"],
+            weight_decay=cfg["weight_decay"],
+        )
+
+    def sup_loss(self, pred, batch):
+        pred = F.softmax(pred, dim=1)
+        loss = self.criterion(pred, batch)
+        loss += self.unsup_loss(batch.x, batch.edge_index, batch.batch)[1]
+        return loss
+
+    def unsup_loss(self, graph_feat, node_feat, batch):
+        local_encode = self.local_dis(node_feat)
+        global_encode = self.global_dis(graph_feat)
+
+        num_graphs = graph_feat.shape[0]
+        num_nodes = node_feat.shape[0]
+
+        pos_mask = torch.zeros((num_nodes, num_graphs)).to(batch.device)
+        neg_mask = torch.ones((num_nodes, num_graphs)).to(batch.device)
+        for nid, gid in enumerate(batch):
+            pos_mask[nid][gid] = 1
+            neg_mask[nid][gid] = 0
+        glob_local_mi = torch.mm(local_encode, global_encode.t())
+        loss = self.mi_loss(pos_mask, neg_mask, glob_local_mi, num_nodes, num_nodes * (num_graphs - 1))
+        return loss
+
+    @staticmethod
+    def mi_loss(pos_mask, neg_mask, mi, pos_div, neg_div):
+        pos_mi = pos_mask * mi
+        neg_mi = neg_mask * mi
+
+        pos_loss = (-math.log(2.0) + F.softplus(-pos_mi)).sum()
+        neg_loss = (-math.log(2.0) + F.softplus(-neg_mi) + neg_mi).sum()
+        # pos_loss = F.softplus(-pos_mi).sum()
+        # neg_loss = (F.softplus(neg_mi)).sum()
+        # pos_loss = pos_mi.sum()
+        # neg_loss = neg_mi.sum()
+        return pos_loss / pos_div + neg_loss / neg_div
+
+
+class FF(nn.Module):
+    r"""Residual MLP layers.
+
+    ..math::
+        out = \mathbf{MLP}(x) + \mathbf{Linear}(x)
+
+    Paramaters
+    ----------
+    in_feats : int
+        Size of each input sample
+    out_feats : int
+        Size of each output sample
+    """
+
+    def __init__(self, in_feats, out_feats):
+        super(FF, self).__init__()
+        self.block = MLP(in_feats, out_feats, out_feats, num_layers=3)
+        self.shortcut = nn.Linear(in_feats, out_feats)
+
+    def forward(self, x):
+        return F.relu(self.block(x)) + self.shortcut(x)
diff --git a/cogdl/wrappers/model_wrapper/heterogeneous/__init__.py b/cogdl/wrappers/model_wrapper/heterogeneous/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/cogdl/wrappers/model_wrapper/heterogeneous/heterogeneous_embedding_mw.py b/cogdl/wrappers/model_wrapper/heterogeneous/heterogeneous_embedding_mw.py
new file mode 100644
index 00000000..c11b6eac
--- /dev/null
+++ b/cogdl/wrappers/model_wrapper/heterogeneous/heterogeneous_embedding_mw.py
@@ -0,0 +1,46 @@
+import argparse
+import numpy as np
+import torch
+
+from sklearn.linear_model import LogisticRegression
+from sklearn.metrics import f1_score
+
+from .. import register_model_wrapper, EmbeddingModelWrapper
+
+
+@register_model_wrapper("heterogeneous_embedding_mw")
+class HeterogeneousEmbeddingModelWrapper(EmbeddingModelWrapper):
+    @staticmethod
+    def add_args(parser: argparse.ArgumentParser):
+        """Add task-specific arguments to the parser."""
+        # fmt: off
+        parser.add_argument("--hidden-size", type=int, default=128)
+        # fmt: on
+
+    def __init__(self, model, hidden_size=200):
+        super(HeterogeneousEmbeddingModelWrapper, self).__init__()
+        self.model = model
+        self.hidden_size = hidden_size
+
+    def train_step(self, batch):
+        embeddings = self.model.train(batch)
+        embeddings = np.hstack((embeddings, batch.x.numpy()))
+
+        return embeddings
+
+    def test_step(self, batch):
+        embeddings, data = batch
+
+        # Select nodes which have label as training data
+        train_index = torch.cat((data.train_node, data.valid_node)).numpy()
+        test_index = data.test_node.numpy()
+        y = data.y.numpy()
+
+        X_train, y_train = embeddings[train_index], y[train_index]
+        X_test, y_test = embeddings[test_index], y[test_index]
+        clf = LogisticRegression()
+        clf.fit(X_train, y_train)
+        preds = clf.predict(X_test)
+        test_f1 = f1_score(y_test, preds, average="micro")
+
+        return dict(f1=test_f1)
diff --git a/cogdl/wrappers/model_wrapper/heterogeneous/heterogeneous_gnn_mw.py b/cogdl/wrappers/model_wrapper/heterogeneous/heterogeneous_gnn_mw.py
new file mode 100644
index 00000000..0557e33d
--- /dev/null
+++ b/cogdl/wrappers/model_wrapper/heterogeneous/heterogeneous_gnn_mw.py
@@ -0,0 +1,43 @@
+import torch
+from cogdl.wrappers.model_wrapper import ModelWrapper, register_model_wrapper
+
+
+@register_model_wrapper("heterogeneous_gnn_mw")
+class HeterogeneousGNNModelWrapper(ModelWrapper):
+    def __init__(self, model, optimizer_config):
+        super(HeterogeneousGNNModelWrapper, self).__init__()
+        self.optimizer_config = optimizer_config
+        self.model = model
+
+    def train_step(self, batch):
+        graph = batch.data
+        pred = self.model(graph)
+        train_mask = graph.train_node
+        loss = self.default_loss_fn(pred[train_mask], graph.y[train_mask])
+        return loss
+
+    def val_step(self, batch):
+        graph = batch.data
+        pred = self.model(graph)
+        val_mask = graph.valid_node
+        loss = self.default_loss_fn(pred[val_mask], graph.y[val_mask])
+        metric = self.evaluate(pred[val_mask], graph.y[val_mask], metric="auto")
+        self.note("val_loss", loss.item())
+        self.note("val_acc", metric)
+
+    def test_step(self, batch):
+        graph = batch.data
+        pred = self.model(graph)
+        test_mask = graph.test_node
+        loss = self.default_loss_fn(pred[test_mask], graph.y[test_mask])
+        metric = self.evaluate(pred[test_mask], graph.y[test_mask], metric="auto")
+        self.note("test_loss", loss.item())
+        self.note("test_acc", metric)
+
+    def setup_optimizer(self):
+        cfg = self.optimizer_config
+        if hasattr(self.model, "get_optimizer"):
+            model_spec_optim = self.model.get_optimizer(cfg)
+            if model_spec_optim is not None:
+                return model_spec_optim
+        return torch.optim.Adam(self.model.parameters(), lr=cfg["lr"], weight_decay=cfg["weight_decay"])
diff --git a/cogdl/wrappers/model_wrapper/heterogeneous/multiplex_embedding_mw.py b/cogdl/wrappers/model_wrapper/heterogeneous/multiplex_embedding_mw.py
new file mode 100644
index 00000000..2ee20911
--- /dev/null
+++ b/cogdl/wrappers/model_wrapper/heterogeneous/multiplex_embedding_mw.py
@@ -0,0 +1,90 @@
+import argparse
+import numpy as np
+import torch
+from sklearn.metrics import auc, f1_score, precision_recall_curve, roc_auc_score
+
+from cogdl.data import Graph
+from .. import register_model_wrapper, EmbeddingModelWrapper
+
+
+def get_score(embs, node1, node2):
+    vector1 = embs[int(node1)]
+    vector2 = embs[int(node2)]
+    return np.dot(vector1, vector2) / (np.linalg.norm(vector1) * np.linalg.norm(vector2))
+
+
+def evaluate(embs, true_edges, false_edges):
+    true_list = list()
+    prediction_list = list()
+    for edge in true_edges:
+        true_list.append(1)
+        prediction_list.append(get_score(embs, edge[0], edge[1]))
+
+    for edge in false_edges:
+        true_list.append(0)
+        prediction_list.append(get_score(embs, edge[0], edge[1]))
+
+    sorted_pred = prediction_list[:]
+    sorted_pred.sort()
+    threshold = sorted_pred[-len(true_edges)]
+
+    y_pred = np.zeros(len(prediction_list), dtype=np.int32)
+    for i in range(len(prediction_list)):
+        if prediction_list[i] >= threshold:
+            y_pred[i] = 1
+
+    y_true = np.array(true_list)
+    y_scores = np.array(prediction_list)
+    ps, rs, _ = precision_recall_curve(y_true, y_scores)
+    return roc_auc_score(y_true, y_scores), f1_score(y_true, y_pred), auc(rs, ps)
+
+
+def evaluate_multiplex(all_embs, test_data):
+    total_roc_auc, total_f1_score, total_pr_auc = [], [], []
+    for key in test_data.keys():
+        embs = all_embs[key]
+        roc_auc, f1_score, pr_auc = evaluate(embs, test_data[key][0], test_data[key][1])
+        total_roc_auc.append(roc_auc)
+        total_f1_score.append(f1_score)
+        total_pr_auc.append(pr_auc)
+    assert len(total_roc_auc) > 0
+    roc_auc, f1_score, pr_auc = (
+        np.mean(total_roc_auc),
+        np.mean(total_f1_score),
+        np.mean(total_pr_auc),
+    )
+    print(f"Test ROC-AUC = {roc_auc:.4f}, F1 = {f1_score:.4f}, PR-AUC = {pr_auc:.4f}")
+    return dict(ROC_AUC=roc_auc, PR_AUC=pr_auc, F1=f1_score)
+
+
+@register_model_wrapper("multiplex_embedding_mw")
+class MultiplexEmbeddingModelWrapper(EmbeddingModelWrapper):
+    @staticmethod
+    def add_args(parser: argparse.ArgumentParser):
+        """Add task-specific arguments to the parser."""
+        # fmt: off
+        parser.add_argument("--hidden-size", type=int, default=200)
+        parser.add_argument("--eval-type", type=str, default='all', nargs='+')
+        # fmt: on
+
+    def __init__(self, model, hidden_size=200, eval_type="all"):
+        super(MultiplexEmbeddingModelWrapper, self).__init__()
+        self.model = model
+        self.hidden_size = hidden_size
+        self.eval_type = eval_type
+
+    def train_step(self, batch):
+        if hasattr(self.model, "multiplicity"):
+            all_embs = self.model(batch)
+        else:
+            all_embs = dict()
+            for key in batch.keys():
+                if self.eval_type == "all" or key in self.eval_type:
+                    G = Graph(edge_index=torch.LongTensor(batch[key]).transpose(0, 1))
+                    embs = self.model(G, return_dict=True)
+                all_embs[key] = embs
+        return all_embs
+
+    def test_step(self, batch):
+        all_embs, test_data = batch
+        return evaluate_multiplex(all_embs, test_data)
diff --git a/cogdl/wrappers/model_wrapper/link_prediction/__init__.py b/cogdl/wrappers/model_wrapper/link_prediction/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/cogdl/wrappers/model_wrapper/link_prediction/embedding_link_prediction_mw.py b/cogdl/wrappers/model_wrapper/link_prediction/embedding_link_prediction_mw.py
new file mode 100644
index 00000000..46eb3ccd
--- /dev/null
+++ b/cogdl/wrappers/model_wrapper/link_prediction/embedding_link_prediction_mw.py
@@ -0,0 +1,53 @@
+import numpy as np
+
+from .. import EmbeddingModelWrapper, register_model_wrapper
+from sklearn.metrics import roc_auc_score, f1_score, auc, precision_recall_curve
+
+
+@register_model_wrapper("embedding_link_prediction_mw")
+class EmbeddingLinkPrediction(EmbeddingModelWrapper):
+    def __init__(self, model):
+        super(EmbeddingLinkPrediction, self).__init__()
+        self.model = model
+
+    def train_step(self, graph):
+        embeddings = self.model(graph)
+        return embeddings
+
+    def test_step(self, batch):
+        embeddings, test_data = batch
+        roc_auc, f1_score, pr_auc = evaluate(embeddings, test_data[0], test_data[1])
+        print(f"Test ROC-AUC = {roc_auc:.4f}, F1 = {f1_score:.4f}, PR-AUC = {pr_auc:.4f}")
+        return dict(ROC_AUC=roc_auc, PR_AUC=pr_auc, F1=f1_score)
+
+
+def evaluate(embs, true_edges, false_edges):
+    true_list = list()
+    prediction_list = list()
+    for edge in true_edges:
+        true_list.append(1)
+        prediction_list.append(get_score(embs, edge[0], edge[1]))
+
+    for edge in false_edges:
+        true_list.append(0)
+        prediction_list.append(get_score(embs, edge[0], edge[1]))
+
+    sorted_pred = prediction_list[:]
+    sorted_pred.sort()
+    threshold = sorted_pred[-len(true_edges)]
+
+    y_pred = np.zeros(len(prediction_list), dtype=np.int32)
+    for i in range(len(prediction_list)):
+        if prediction_list[i] >= threshold:
+            y_pred[i] = 1
+
+    y_true = np.array(true_list)
+    y_scores = np.array(prediction_list)
+    ps, rs, _ = precision_recall_curve(y_true, y_scores)
+    return roc_auc_score(y_true, y_scores), f1_score(y_true, y_pred), auc(rs, ps)
+
+
+def get_score(embs, node1, node2, eps=1e-5):
+    vector1 = embs[int(node1)]
+    vector2 = embs[int(node2)]
+    return np.dot(vector1, vector2) / (np.linalg.norm(vector1) * np.linalg.norm(vector2) + eps)
diff --git a/cogdl/wrappers/model_wrapper/link_prediction/gnn_kg_link_prediction_mw.py b/cogdl/wrappers/model_wrapper/link_prediction/gnn_kg_link_prediction_mw.py
new file mode 100644
index 00000000..1b9c43dc
--- /dev/null
+++ b/cogdl/wrappers/model_wrapper/link_prediction/gnn_kg_link_prediction_mw.py
@@ -0,0 +1,77 @@
+import torch
+import torch.nn as nn
+
+from .. import register_model_wrapper, ModelWrapper
+from cogdl.utils.link_prediction_utils import cal_mrr, DistMultLayer, ConvELayer
+
+
+@register_model_wrapper("gnn_kg_link_prediction_mw")
+class GNNKGLinkPrediction(ModelWrapper):
+    @staticmethod
+    def add_args(parser):
+        # fmt: off
+        parser.add_argument("--score-func", type=str, default="distmult")
+        # fmt: on
+
+    def __init__(self, model, optimizer_cfg, score_func):
+        super(GNNKGLinkPrediction, self).__init__()
+
+        self.model = model
+        self.optimizer_cfg = optimizer_cfg
+        hidden_size = optimizer_cfg["hidden_size"]
+
+        self.score_func = score_func
+        if score_func == "distmult":
+            self.scoring = DistMultLayer()
+        elif score_func == "conve":
+            self.scoring = ConvELayer(hidden_size)
+        else:
+            raise NotImplementedError
+
+    def train_step(self, subgraph):
+        graph = subgraph
+        mask = graph.train_mask
+        edge_index = torch.stack(graph.edge_index)
+        edge_index, edge_types = edge_index[:, mask], graph.edge_attr[mask]
+
+        with graph.local_graph():
+            graph.edge_index = edge_index
+            graph.edge_attr = edge_types
+            loss = self.model.loss(graph)
+        return loss
+
+    def val_step(self, subgraph):
+        train_mask = subgraph.train_mask
+        eval_mask = subgraph.val_mask
+        return self.eval_step(subgraph, train_mask, eval_mask)
+
+    def test_step(self, subgraph):
+        infer_mask = subgraph.train_mask | subgraph.val_mask
+        eval_mask = subgraph.test_mask
+        return self.eval_step(subgraph, infer_mask, eval_mask)
+
+    def eval_step(self, graph, mask1, mask2):
+        row, col = graph.edge_index
+        edge_types = graph.edge_attr
+
+        with graph.local_graph():
+            graph.edge_index = (row[mask1], col[mask1])
+            graph.edge_attr = edge_types[mask1]
+            output, rel_weight = self.model.predict(graph)
+
+        mrr, hits = cal_mrr(
+            output,
+            rel_weight,
+            (row[mask2], col[mask2]),
+            edge_types[mask2],
+            scoring=self.scoring,
+            protocol="raw",
+            batch_size=500,
+            hits=[1, 3, 10],
+        )
+
+        return dict(mrr=mrr, hits1=hits[0], hits3=hits[1], hits10=hits[2])
+
+    def setup_optimizer(self):
+        lr, weight_decay = self.optimizer_cfg["lr"], self.optimizer_cfg["weight_decay"]
+        return torch.optim.AdamW(self.parameters(), lr=lr, weight_decay=weight_decay)
diff --git a/cogdl/wrappers/model_wrapper/link_prediction/gnn_link_prediction_mw.py b/cogdl/wrappers/model_wrapper/link_prediction/gnn_link_prediction_mw.py
new file mode 100644
index 00000000..4b479525
--- /dev/null
+++ b/cogdl/wrappers/model_wrapper/link_prediction/gnn_link_prediction_mw.py
@@ -0,0 +1,61 @@
+from sklearn.metrics import roc_auc_score
+
+import torch
+from .. import ModelWrapper, register_model_wrapper
+from cogdl.utils import negative_edge_sampling
+
+
+@register_model_wrapper("gnn_link_prediction_mw")
+class GNNLinkPredictionModelWrapper(ModelWrapper):
+    def __init__(self, model, optimizer_cfg):
+        super(GNNLinkPredictionModelWrapper, self).__init__()
+        self.model = model
+        self.optimizer_cfg = optimizer_cfg
+        self.loss_fn = torch.nn.BCELoss()
+
+    def train_step(self, subgraph):
+        graph = subgraph
+
+        train_neg_edges = negative_edge_sampling(graph.train_edges, graph.num_nodes).to(self.device)
+        train_pos_edges = graph.train_edges
+        edge_index = torch.cat([train_pos_edges, train_neg_edges], dim=1)
+        labels = self.get_link_labels(train_pos_edges.shape[1], train_neg_edges.shape[1], self.device)
+
+        # link prediction loss
+        with graph.local_graph():
+            graph.edge_index = edge_index
+            emb = self.model(graph)
+        pred = (emb[edge_index[0]] * emb[edge_index[1]]).sum(1)
+        pred = torch.sigmoid(pred)
+        loss = self.loss_fn(pred, labels)
+        return loss
+
+    def val_step(self, subgraph):
+        graph = subgraph
+        pos_edges = graph.val_edges
+        neg_edges = graph.val_neg_edges
+        train_edges = graph.train_edges
+        edges = torch.cat([pos_edges, neg_edges], dim=1)
+        labels = self.get_link_labels(pos_edges.shape[1], neg_edges.shape[1], self.device).long()
+        with graph.local_graph():
+            graph.edge_index = train_edges
+            with torch.no_grad():
+                emb = self.model(graph)
+                pred = (emb[edges[0]] * emb[edges[1]]).sum(-1)
+        pred = torch.sigmoid(pred)
+
+        auc_score = roc_auc_score(labels.cpu().numpy(), pred.cpu().numpy())
+
+        self.note("auc", auc_score)
+
+    @staticmethod
+    def get_link_labels(num_pos, num_neg, device=None):
+        labels = torch.zeros(num_pos + num_neg)
+        labels[:num_pos] = 1
+        if device is not None:
+            labels = labels.to(device)
+        return labels.float()
+
+    def setup_optimizer(self):
+        lr, wd = self.optimizer_cfg["lr"], self.optimizer_cfg["weight_decay"]
+        return torch.optim.Adam(self.parameters(), lr=lr, weight_decay=wd)
diff --git a/cogdl/wrappers/model_wrapper/node_classification/__init__.py b/cogdl/wrappers/model_wrapper/node_classification/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/cogdl/wrappers/model_wrapper/node_classification/dgi_mw.py b/cogdl/wrappers/model_wrapper/node_classification/dgi_mw.py
new file mode 100644
index 00000000..6362c326
--- /dev/null
+++ b/cogdl/wrappers/model_wrapper/node_classification/dgi_mw.py
@@ -0,0 +1,117 @@
+import numpy as np
+
+import torch
+import torch.nn as nn
+
+from .. import register_model_wrapper, ModelWrapper
+from cogdl.wrappers.tools.wrapper_utils import evaluate_node_embeddings_using_logreg
+
+
+@register_model_wrapper("dgi_mw")
+class DGIModelWrapper(ModelWrapper):
+    @staticmethod
+    def add_args(parser):
+        # fmt: off
+        parser.add_argument("--hidden-size", type=int, default=512)
+        # fmt: on
+
+    def __init__(self, model, optimizer_cfg):
+        super(DGIModelWrapper, self).__init__()
+        self.model = model
+        self.optimizer_cfg = optimizer_cfg
+        self.read = AvgReadout()
+        self.sigm = nn.Sigmoid()
+
+        hidden_size = optimizer_cfg["hidden_size"]
+        assert hidden_size > 0
+        self.disc = Discriminator(hidden_size)
+        self.loss_fn = torch.nn.BCEWithLogitsLoss()
+        self.act = nn.PReLU()
+
+    def train_step(self, subgraph):
+        graph = subgraph
+        graph.sym_norm()
+        x = graph.x
+        shuffle_x = self.augment(graph)
+
+        graph.x = x
+        h_pos = self.act(self.model(graph))
+        c = self.read(h_pos)
+        c = self.sigm(c)
+
+        graph.x = shuffle_x
+        h_neg = self.act(self.model(graph))
+        logits = self.disc(c, h_pos, h_neg)
+        graph.x = x
+
+        num_nodes = x.shape[0]
+        labels = torch.zeros((num_nodes * 2,), device=x.device)
+        labels[:num_nodes] = 1
+        loss = self.loss_fn(logits, labels)
+        return loss
+
+    # def evaluate(self, dataset):
+    #     device = self.device
+    #     graph = dataset.data.to(device)
+
+    def test_step(self, graph):
+        with torch.no_grad():
+            pred = self.act(self.model(graph))
+        y = graph.y
+        result = evaluate_node_embeddings_using_logreg(pred, y, graph.train_mask, graph.test_mask)
+        self.note("test_microf1", result)
+
+    @staticmethod
+    def augment(graph):
+        idx = np.random.permutation(graph.num_nodes)
+        shuffle_x = graph.x[idx, :]
+        return shuffle_x
+
+    def setup_optimizer(self):
+        cfg = self.optimizer_cfg
+        return torch.optim.Adam(self.parameters(), lr=cfg["lr"], weight_decay=cfg["weight_decay"])
+
+
+# Borrowed from https://github.com/PetarV-/DGI
+class AvgReadout(nn.Module):
+    def __init__(self):
+        super(AvgReadout, self).__init__()
+
+    def forward(self, seq, msk=None):
+        dim = len(seq.shape) - 2
+        if msk is None:
+            return torch.mean(seq, dim)
+        else:
+            return torch.sum(seq * msk, dim) / torch.sum(msk)
+
+
+# Borrowed from https://github.com/PetarV-/DGI
+class Discriminator(nn.Module):
+    def __init__(self, n_h):
+        super(Discriminator, self).__init__()
+        self.f_k = nn.Bilinear(n_h, n_h, 1)
+
+        for m in self.modules():
+            self.weights_init(m)
+
+    def weights_init(self, m):
+        if isinstance(m, nn.Bilinear):
+            torch.nn.init.xavier_uniform_(m.weight.data)
+            if m.bias is not None:
+                m.bias.data.fill_(0.0)
+
+    def forward(self, c, h_pl, h_mi, s_bias1=None, s_bias2=None):
+        c_x = torch.unsqueeze(c, 0)
+        c_x = c_x.expand_as(h_pl)
+
+        sc_1 = torch.squeeze(self.f_k(h_pl, c_x), 1)
+        sc_2 = torch.squeeze(self.f_k(h_mi, c_x), 1)
+
+        if s_bias1 is not None:
+            sc_1 += s_bias1
+        if s_bias2 is not None:
+            sc_2 += s_bias2
+
+        logits = torch.cat((sc_1, sc_2))
+
+        return logits
diff --git a/cogdl/wrappers/model_wrapper/node_classification/gcnmix_mw.py b/cogdl/wrappers/model_wrapper/node_classification/gcnmix_mw.py
new file mode 100644
index 00000000..d596778a
--- /dev/null
+++ b/cogdl/wrappers/model_wrapper/node_classification/gcnmix_mw.py
@@ -0,0 +1,166 @@
+import copy
+import random
+import numpy as np
+import torch
+
+from .. import register_model_wrapper, ModelWrapper
+from cogdl.wrappers.tools.wrapper_utils import pre_evaluation_index
+
+
+@register_model_wrapper("gcnmix_mw")
+class GCNMixModelWrapper(ModelWrapper):
+    """
+    GCNMixModelWrapper calls `forward_aux` in model
+    `forward_aux` is similar to `forward` but ignores `spmm` operation.
+    """
+
+    @staticmethod
+    def add_args(parser):
+        # fmt: off
+        parser.add_argument("--temperature", type=float, default=0.1)
+        parser.add_argument("--rampup-starts", type=int, default=500)
+        parser.add_argument("--rampup-ends", type=int, default=1000)
+        parser.add_argument("--mixup-consistency", type=float, default=10.0)
+        parser.add_argument("--ema-decay", type=float, default=0.999)
+        parser.add_argument("--tau", type=float, default=1.0)
+        parser.add_argument("--k", type=int, default=10)
+        # fmt: on
+
+    def __init__(
+        self, model, optimizer_cfg, temperature, rampup_starts, rampup_ends, mixup_consistency, ema_decay, tau, k
+    ):
+        super(GCNMixModelWrapper, self).__init__()
+        self.optimizer_cfg = optimizer_cfg
+        self.temperature = temperature
+        self.ema_decay = ema_decay
+        self.tau = tau
+        self.k = k
+
+        self.model = model
+        self.model_ema = copy.deepcopy(self.model)
+
+        for p in self.model_ema.parameters():
+            p.detach_()
+        self.epoch = 0
+        self.opt = {
+            "epoch": 0,
+            "final_consistency_weight": mixup_consistency,
+            "rampup_starts": rampup_starts,
+            "rampup_ends": rampup_ends,
+        }
+        self.mix_loss = torch.nn.BCELoss()
+        self.mix_transform = None
+
+    def train_step(self, subgraph):
+        if self.mix_transform is None:
+            if len(subgraph.y.shape) > 1:
+                self.mix_transform = torch.nn.Sigmoid()
+            else:
+                self.mix_transform = torch.nn.Softmax(-1)
+        graph = subgraph
+        device = graph.x.device
+        train_mask = graph.train_mask
+
+        self.opt["epoch"] += 1
+
+        rand_n = random.randint(0, 1)
+        if rand_n == 0:
+            vector_labels = get_one_hot_label(graph.y, train_mask).to(device)
+            loss = self.update_aux(graph, vector_labels, train_mask)
+        else:
+            loss = self.update_soft(graph)
+
+        alpha = min(1 - 1 / (self.epoch + 1), self.ema_decay)
+        for ema_param, param in zip(self.model_ema.parameters(), self.model.parameters()):
+            ema_param.data.mul_(alpha).add_((1 - alpha) * param.data)
+
+        return loss
+
+    def val_step(self, subgraph):
+        graph = subgraph
+        val_mask = graph.val_mask
+        pred = self.model_ema(graph)
+        loss = self.default_loss_fn(pred[val_mask], graph.y[val_mask])
+
+        metric = self.evaluate(pred[val_mask], graph.y[val_mask], metric="auto")
+
+        self.note("val_loss", loss.item())
+        self.note("val_metric", metric)
+
+    def test_step(self, subgraph):
+        test_mask = subgraph.test_mask
+        pred = self.model_ema(subgraph)
+        loss = self.default_loss_fn(pred[test_mask], subgraph.y[test_mask])
+
+        metric = self.evaluate(pred[test_mask], subgraph.y[test_mask], metric="auto")
+
+        self.note("test_loss", loss.item())
+        self.note("test_metric", metric)
+
+    def update_soft(self, graph):
+        out = self.model(graph)
+        train_mask = graph.train_mask
+        loss_sum = self.default_loss_fn(out[train_mask], graph.y[train_mask])
+        return loss_sum
+
+    def update_aux(self, data, vector_labels, train_index):
+        device = self.device
+        train_unlabelled = torch.where(~data.train_mask)[0].to(device)
+        temp_labels = torch.zeros(self.k, vector_labels.shape[0], vector_labels.shape[1]).to(device)
+        with torch.no_grad():
+            for i in range(self.k):
+                temp_labels[i, :, :] = self.model(data) / self.tau
+
+        target_labels = temp_labels.mean(dim=0)
+        target_labels = sharpen(target_labels, self.temperature)
+        vector_labels[train_unlabelled] = target_labels[train_unlabelled]
+        sampled_unlabelled = torch.randint(0, train_unlabelled.shape[0], size=(train_index.shape[0],))
+        train_unlabelled = train_unlabelled[sampled_unlabelled]
+
+        def get_loss(index):
+            # TODO: call `forward_aux` in model
+            mix_logits, target = self.model.forward_aux(data.x, vector_labels, index, mix_hidden=True)
+            # temp_loss = self.loss_f(F.softmax(mix_logits[index], -1), target)
+            temp_loss = self.mix_loss(self.mix_transform(mix_logits[index]), target)
+            return temp_loss
+
+        sup_loss = get_loss(train_index)
+        unsup_loss = get_loss(train_unlabelled)
+
+        mixup_weight = get_current_consistency_weight(
+            self.opt["final_consistency_weight"], self.opt["rampup_starts"], self.opt["rampup_ends"], self.opt["epoch"]
+        )
+
+        loss_sum = sup_loss + mixup_weight * unsup_loss
+        return loss_sum
+
+    def setup_optimizer(self):
+        lr = self.optimizer_cfg["lr"]
+        wd = self.optimizer_cfg["weight_decay"]
+        return torch.optim.Adam(self.parameters(), lr=lr, weight_decay=wd)
+
+
+def get_one_hot_label(labels, index):
+    num_classes = int(torch.max(labels) + 1)
+    target = torch.zeros(labels.shape[0], num_classes).to(labels.device)
+
+    target[index, labels[index]] = 1
+    return target
+
+
+def sharpen(prob, temperature):
+    prob = torch.pow(prob, 1.0 / temperature)
+    row_sum = torch.sum(prob, dim=1).reshape(-1, 1)
+    return prob / row_sum
+
+
+def get_current_consistency_weight(final_consistency_weight, rampup_starts, rampup_ends, epoch):
+    # Consistency ramp-up from https://arxiv.org/abs/1610.02242
+    rampup_length = rampup_ends - rampup_starts
+    rampup = 1.0
+    epoch = epoch - rampup_starts
+    if rampup_length != 0:
+        current = np.clip(epoch, 0.0, rampup_length)
+        phase = 1.0 - current / rampup_length
+        rampup = float(np.exp(-5.0 * phase * phase))
+    return final_consistency_weight * rampup
diff --git a/cogdl/wrappers/model_wrapper/node_classification/grace_mw.py b/cogdl/wrappers/model_wrapper/node_classification/grace_mw.py
new file mode 100644
index 00000000..827da6d0
--- /dev/null
+++ b/cogdl/wrappers/model_wrapper/node_classification/grace_mw.py
@@ -0,0 +1,99 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from cogdl.data import Graph
+from .. import register_model_wrapper, ModelWrapper
+from cogdl.wrappers.tools.wrapper_utils import evaluate_node_embeddings_using_logreg
+from cogdl.utils import dropout_adj, dropout_features
+
+
+@register_model_wrapper("grace_mw")
+class GRACEModelWrapper(ModelWrapper):
+    @staticmethod
+    def add_args(parser):
+        # fmt: off
+        parser.add_argument("--tau", type=float, default=0.5)
+        parser.add_argument("--drop-feature-rates", type=float, nargs="+", default=[0.3, 0.4])
+        parser.add_argument("--drop-edge-rates", type=float, nargs="+", default=[0.2, 0.4])
+        parser.add_argument("--batch-fwd", type=int, default=-1)
+        parser.add_argument("--proj-hidden-size", type=int, default=128)
+        # fmt: on
+
+    def __init__(self, model, optimizer_cfg, tau, drop_feature_rates, drop_edge_rates, batch_fwd, proj_hidden_size):
+        super(GRACEModelWrapper, self).__init__()
+        self.tau = tau
+        self.drop_feature_rates = drop_feature_rates
+        self.drop_edge_rates = drop_edge_rates
+        self.batch_size = batch_fwd
+
+        self.model = model
+        hidden_size = optimizer_cfg["hidden_size"]
+        self.project_head = nn.Sequential(
+            nn.Linear(hidden_size, proj_hidden_size), nn.ELU(), nn.Linear(proj_hidden_size, hidden_size)
+        )
+        self.optimizer_cfg = optimizer_cfg
+
+    def train_step(self, subgraph):
+        graph = subgraph
+        z1 = self.prop(graph, graph.x, self.drop_feature_rates[0], self.drop_edge_rates[0])
+        z2 = self.prop(graph, graph.x, self.drop_feature_rates[1], self.drop_edge_rates[1])
+
+        z1 = self.project_head(z1)
+        z2 = self.project_head(z2)
+
+        if self.batch_size > 0:
+            return 0.5 * (self.batched_loss(z1, z2, self.batch_size) + self.batched_loss(z2, z1, self.batch_size))
+        else:
+            return 0.5 * (self.contrastive_loss(z1, z2) + self.contrastive_loss(z2, z1))
+
+    def test_step(self, graph):
+        with torch.no_grad():
+            pred = self.model(graph)
+        y = graph.y
+        result = evaluate_node_embeddings_using_logreg(pred, y, graph.train_mask, graph.test_mask)
+        self.note("acc", result)
+
+    def prop(
+        self,
+        graph: Graph,
+        x: torch.Tensor,
+        drop_feature_rate: float = 0.0,
+        drop_edge_rate: float = 0.0,
+    ):
+        x = dropout_features(x, drop_feature_rate)
+        with graph.local_graph():
+            graph = dropout_adj(graph.edge_index, graph.edge_weight, drop_edge_rate)
+            return self.model.forward(graph, x)
+
+    def contrastive_loss(self, z1: torch.Tensor, z2: torch.Tensor):
+        z1 = F.normalize(z1, p=2, dim=-1)
+        z2 = F.normalize(z2, p=2, dim=-1)
+
+        def score_func(emb1, emb2):
+            scores = torch.matmul(emb1, emb2.t())
+            scores = torch.exp(scores / self.tau)
+            return scores
+
+        intro_scores = score_func(z1, z1)
+        inter_scores = score_func(z1, z2)
+
+        _loss = -torch.log(intro_scores.diag() / (intro_scores.sum(1) - intro_scores.diag() + inter_scores.sum(1)))
+        return torch.mean(_loss)
+
+    def batched_loss(
+        self,
+        z1: torch.Tensor,
+        z2: torch.Tensor,
+        batch_size: int,
+    ):
+        num_nodes = z1.shape[0]
+        num_batches = (num_nodes - 1) // batch_size + 1
+
+        losses = []
+        indices = torch.arange(num_nodes).to(z1.device)
+        for i in range(num_batches):
+            train_indices = indices[i * batch_size : (i + 1) * batch_size]
+            _loss = self.contrastive_loss(z1[train_indices], z2)
+            losses.append(_loss)
+        return sum(losses) / len(losses)
diff --git a/cogdl/wrappers/model_wrapper/node_classification/grand_mw.py b/cogdl/wrappers/model_wrapper/node_classification/grand_mw.py
new file mode 100644
index 00000000..f333d52f
--- /dev/null
+++ b/cogdl/wrappers/model_wrapper/node_classification/grand_mw.py
@@ -0,0 +1,64 @@
+import torch
+import torch.nn.functional as F
+
+from cogdl.wrappers.model_wrapper import register_model_wrapper
+from cogdl.wrappers.model_wrapper.node_classification.node_classification_mw import NodeClfModelWrapper
+
+
+@register_model_wrapper("grand_mw")
+class GrandModelWrapper(NodeClfModelWrapper):
+    """
+    sample : int
+        Number of augmentations for consistency loss
+    temperature : float
+        Temperature to sharpen predictions.
+    lmbda : float
+         Proportion of consistency loss of unlabelled data
+    """
+
+    @staticmethod
+    def add_args(parser):
+        # fmt: off
+        parser.add_argument("--temperature", type=float, default=0.5)
+        parser.add_argument("--lmbda", type=float, default=0.5)
+        parser.add_argument("--sample", type=int, default=2)
+        # fmt: on
+
+    def __init__(self, model, optimizer_cfg, sample=2, temperature=0.5, lmbda=0.5):
+        super(GrandModelWrapper, self).__init__(model, optimizer_cfg)
+        self.sample = sample
+        self.temperature = temperature
+        self.lmbda = lmbda
+
+    def train_step(self, batch):
+        graph = batch
+        output_list = []
+        for i in range(self.sample):
+            output_list.append(self.model(graph))
+        loss_train = 0.0
+        for output in output_list:
+            loss_train += self.default_loss_fn(output[graph.train_mask], graph.y[graph.train_mask])
+        loss_train = loss_train / self.sample
+
+        if len(graph.y.shape) > 1:
+            output_list = [torch.sigmoid(x) for x in output_list]
+        else:
+            output_list = [F.log_softmax(x, dim=-1) for x in output_list]
+        loss_consis = self.consistency_loss(output_list, graph.train_mask)
+
+        return loss_train + loss_consis
+
+    def consistency_loss(self, logps, train_mask):
+        temp = self.temperature
+        ps = [torch.exp(p)[~train_mask] for p in logps]
+        sum_p = 0.0
+        for p in ps:
+            sum_p = sum_p + p
+        avg_p = sum_p / len(ps)
+        sharp_p = (torch.pow(avg_p, 1.0 / temp) / torch.sum(torch.pow(avg_p, 1.0 / temp), dim=1, keepdim=True)).detach()
+        loss = 0.0
+        for p in ps:
+            loss += torch.mean((p - sharp_p).pow(2).sum(1))
+        loss = loss / len(ps)
+
+        return self.lmbda * loss
diff --git a/cogdl/wrappers/model_wrapper/node_classification/graphsage_mw.py b/cogdl/wrappers/model_wrapper/node_classification/graphsage_mw.py
new file mode 100644
index 00000000..cefddc24
--- /dev/null
+++ b/cogdl/wrappers/model_wrapper/node_classification/graphsage_mw.py
@@ -0,0 +1,46 @@
+import torch
+
+from .. import ModelWrapper, register_model_wrapper
+from cogdl.wrappers.tools.wrapper_utils import pre_evaluation_index
+
+
+@register_model_wrapper("graphsage_mw")
+class GraphSAGEModelWrapper(ModelWrapper):
+    def __init__(self, model, optimizer_cfg):
+        super(GraphSAGEModelWrapper, self).__init__()
+        self.model = model
+        self.optimizer_cfg = optimizer_cfg
+
+    def train_step(self, batch):
+        x_src, y, adjs = batch
+        pred = self.model(x_src, adjs)
+        loss = self.default_loss_fn(pred, y)
+        return loss
+
+    def val_step(self, batch):
+        x_src, y, adjs = batch
+        pred = self.model(x_src, adjs)
+        loss = self.default_loss_fn(pred, y)
+
+        metric = self.evaluate(pred, y, metric="auto")
+
+        self.note("val_loss", loss.item())
+        self.note("val_metric", metric)
+
+    def test_step(self, batch):
+        dataset, test_loader = batch
+        graph = dataset.data
+        if hasattr(self.model, "inference"):
+            pred = self.model.inference(dataset.graph.x, test_loader)
+        else:
+            pred = self.model(graph)
+        pred = pred[graph.test_mask]
+        y = graph.y[graph.test_mask]
+
+        metric = self.evaluate(pred, y, metric="auto")
+        self.note("test_loss", self.default_loss_fn(pred, y))
+        self.note("test_metric", metric)
+
+    def setup_optimizer(self):
+        cfg = self.optimizer_cfg
+        return torch.optim.Adam(self.parameters(), lr=cfg["lr"], weight_decay=cfg["weight_decay"])
diff --git a/cogdl/wrappers/model_wrapper/node_classification/m3s_mw.py b/cogdl/wrappers/model_wrapper/node_classification/m3s_mw.py
new file mode 100644
index 00000000..3f9cccd1
--- /dev/null
+++ b/cogdl/wrappers/model_wrapper/node_classification/m3s_mw.py
@@ -0,0 +1,77 @@
+import copy
+
+import torch
+import numpy as np
+from sklearn.cluster import KMeans
+
+from cogdl.wrappers.data_wrapper import DataWrapper
+
+from .. import register_model_wrapper
+from .node_classification_mw import NodeClfModelWrapper
+
+
+@register_model_wrapper("m3s_mw")
+class M3SModelWrapper(NodeClfModelWrapper):
+    @staticmethod
+    def add_args(parser):
+        # fmt: off
+        parser.add_argument("--n-cluster", type=int, default=10)
+        parser.add_argument("--num-new-labels", type=int, default=10)
+        # fmt: on
+
+    def __init__(self, model, optimizer_cfg, n_cluster, num_new_labels):
+        super(M3SModelWrapper, self).__init__(model, optimizer_cfg)
+        self.model = model
+        self.num_clusters = n_cluster
+        self.hidden_size = optimizer_cfg["hidden_size"]
+        self.num_new_labels = num_new_labels
+        self.optimizer_cfg = optimizer_cfg
+
+    def pre_stage(self, stage, data_w: DataWrapper):
+        if stage > 0:
+            graph = data_w.get_dataset().data
+            graph.store("y")
+
+            y = copy.deepcopy(graph.y)
+
+            num_classes = graph.num_classes
+            num_nodes = graph.num_nodes
+
+            with torch.no_grad():
+                emb = self.model.get_embeddings(graph)
+
+            confidence_ranking = np.zeros([num_classes, num_nodes], dtype=int)
+            kmeans = KMeans(n_clusters=self.num_clusters, random_state=0).fit(emb)
+            clusters = kmeans.labels_
+
+            # Compute centroids μ_m of each class m in labeled data and v_l of each cluster l in unlabeled data.
+            labeled_centroid = np.zeros([num_classes, self.hidden_size])
+            unlabeled_centroid = np.zeros([self.num_clusters, self.hidden_size])
+            for i in range(num_nodes):
+                if graph.train_mask[i]:
+                    labeled_centroid[y[i]] += emb[i]
+                else:
+                    unlabeled_centroid[clusters[i]] += emb[i]
+
+            # Align labels for each cluster
+            align = np.zeros(self.num_clusters, dtype=int)
+            for i in range(self.num_clusters):
+                for j in range(num_classes):
+                    if np.linalg.norm(unlabeled_centroid[i] - labeled_centroid[j]) < np.linalg.norm(
+                        unlabeled_centroid[i] - labeled_centroid[align[i]]
+                    ):
+                        align[i] = j
+
+            # Add new labels
+            for i in range(num_classes):
+                t = self.num_new_labels
+                for j in range(num_nodes):
+                    idx = confidence_ranking[i][j]
+                    if not graph.train_mask[idx]:
+                        if t <= 0:
+                            break
+                        t -= 1
+                        if align[clusters[idx]] == i:
+                            graph.train_mask[idx] = True
+                            y[idx] = i
+            return y
diff --git a/cogdl/wrappers/model_wrapper/node_classification/mvgrl_mw.py b/cogdl/wrappers/model_wrapper/node_classification/mvgrl_mw.py
new file mode 100644
index 00000000..60ee5261
--- /dev/null
+++ b/cogdl/wrappers/model_wrapper/node_classification/mvgrl_mw.py
@@ -0,0 +1,37 @@
+import torch
+import torch.nn as nn
+
+from .. import register_model_wrapper, ModelWrapper
+from cogdl.wrappers.tools.wrapper_utils import evaluate_node_embeddings_using_logreg
+
+
+@register_model_wrapper("mvgrl_mw")
+class MVGRLModelWrapper(ModelWrapper):
+    def __init__(self, model, optimizer_cfg):
+        super(MVGRLModelWrapper, self).__init__()
+        self.model = model
+        self.optimizer_cfg = optimizer_cfg
+        self.loss_f = nn.BCEWithLogitsLoss()
+
+    def train_step(self, subgraph):
+        graph = subgraph
+        logits = self.model(graph)
+        labels = torch.zeros_like(logits)
+        num_outs = logits.shape[1]
+        labels[:, : num_outs // 2] = 1
+        loss = self.loss_f(logits, labels)
+        return loss
+
+    # def evaluate(self, dataset):
+    #     self.model.eval()
+    # graph = dataset.data
+    def test_step(self, graph):
+        with torch.no_grad():
+            pred = self.model(graph)
+        y = graph.y
+        result = evaluate_node_embeddings_using_logreg(pred, y, graph.train_mask, graph.test_mask)
+        self.note("test_acc", result)
+
+    def setup_optimizer(self):
+        cfg = self.optimizer_cfg
+        return torch.optim.Adam(self.parameters(), lr=cfg["lr"], weight_decay=cfg["weight_decay"])
diff --git a/cogdl/wrappers/model_wrapper/node_classification/network_embedding_mw.py b/cogdl/wrappers/model_wrapper/node_classification/network_embedding_mw.py
new file mode 100644
index 00000000..d51a4435
--- /dev/null
+++ b/cogdl/wrappers/model_wrapper/node_classification/network_embedding_mw.py
@@ -0,0 +1,27 @@
+from cogdl.wrappers.tools.wrapper_utils import evaluate_node_embeddings_using_liblinear
+from .. import register_model_wrapper, EmbeddingModelWrapper
+
+
+@register_model_wrapper("network_embedding_mw")
+class NetworkEmbeddingModelWrapper(EmbeddingModelWrapper):
+    @staticmethod
+    def add_args(parser):
+        # fmt: off
+        parser.add_argument("--num-shuffle", type=int, default=10)
+        parser.add_argument("--training-percents", default=[0.9], type=float, nargs="+")
+        # parser.add_argument("--enhance", type=str, default=None, help="use prone or prone++ to enhance embedding")
+        # fmt: on
+
+    def __init__(self, model, num_shuffle=1, training_percents=[0.1]):
+        super(NetworkEmbeddingModelWrapper, self).__init__()
+        self.model = model
+        self.num_shuffle = num_shuffle
+        self.training_percents = training_percents
+
+    def train_step(self, batch):
+        emb = self.model(batch)
+        return emb
+
+    def test_step(self, batch):
+        x, y = batch
+        return evaluate_node_embeddings_using_liblinear(x, y, self.num_shuffle, self.training_percents)
diff --git a/cogdl/wrappers/model_wrapper/node_classification/node_classification_mw.py b/cogdl/wrappers/model_wrapper/node_classification/node_classification_mw.py
new file mode 100644
index 00000000..c480cd36
--- /dev/null
+++ b/cogdl/wrappers/model_wrapper/node_classification/node_classification_mw.py
@@ -0,0 +1,51 @@
+import torch
+from cogdl.wrappers.model_wrapper import ModelWrapper, register_model_wrapper
+
+
+@register_model_wrapper("node_classification_mw")
+class NodeClfModelWrapper(ModelWrapper):
+    def __init__(self, model, optimizer_config):
+        super(NodeClfModelWrapper, self).__init__()
+        self.optimizer_config = optimizer_config
+        self.model = model
+
+    def train_step(self, subgraph):
+        graph = subgraph
+        pred = self.model(graph)
+        train_mask = graph.train_mask
+        loss = self.default_loss_fn(pred[train_mask], graph.y[train_mask])
+        return loss
+
+    def val_step(self, subgraph):
+        graph = subgraph
+        pred = self.model(graph)
+        y = graph.y
+        val_mask = graph.val_mask
+        loss = self.default_loss_fn(pred[val_mask], y[val_mask])
+
+        metric = self.evaluate(pred[val_mask], graph.y[val_mask], metric="auto")
+
+        self.note("val_loss", loss.item())
+        self.note("val_metric", metric)
+
+    def test_step(self, batch):
+        graph = batch
+        pred = self.model(graph)
+        test_mask = batch.test_mask
+        loss = self.default_loss_fn(pred[test_mask], batch.y[test_mask])
+
+        metric = self.evaluate(pred[test_mask], batch.y[test_mask], metric="auto")
+
+        self.note("test_loss", loss.item())
+        self.note("test_metric", metric)
+
+    def setup_optimizer(self):
+        cfg = self.optimizer_config
+        if hasattr(self.model, "setup_optimizer"):
+            model_spec_optim = self.model.setup_optimizer(cfg)
+            if model_spec_optim is not None:
+                return model_spec_optim
+        return torch.optim.Adam(self.model.parameters(), lr=cfg["lr"], weight_decay=cfg["weight_decay"])
+
+    def set_early_stopping(self):
+        return "val_metric", ">"
diff --git a/cogdl/wrappers/model_wrapper/node_classification/pprgo_mw.py b/cogdl/wrappers/model_wrapper/node_classification/pprgo_mw.py
new file mode 100644
index 00000000..31e51b5b
--- /dev/null
+++ b/cogdl/wrappers/model_wrapper/node_classification/pprgo_mw.py
@@ -0,0 +1,46 @@
+import torch
+from cogdl.wrappers.model_wrapper import ModelWrapper, register_model_wrapper
+from cogdl.wrappers.tools.wrapper_utils import pre_evaluation_index
+
+
+@register_model_wrapper("pprgo_mw")
+class PPRGoModelWrapper(ModelWrapper):
+    def __init__(self, model, optimizer_config):
+        super(PPRGoModelWrapper, self).__init__()
+        self.optimizer_config = optimizer_config
+        self.model = model
+
+    def train_step(self, batch):
+        x, targets, ppr_scores, y = batch
+        pred = self.model(x, targets, ppr_scores)
+        loss = self.default_loss_fn(pred, y)
+        return loss
+
+    def val_step(self, batch):
+        graph = batch
+        pred = self.model.predict(graph)
+
+        y = graph.y[graph.val_mask]
+        pred = pred[graph.val_mask]
+
+        loss = self.default_loss_fn(pred, y)
+
+        metric = self.evaluate(pred, y, metric="auto")
+
+        self.note("val_loss", loss.item())
+        self.note("val_metric", metric)
+
+    def test_step(self, batch):
+        graph = batch
+        pred = self.model.predict(graph)
+        test_mask = batch.test_mask
+        pred = pred[test_mask]
+        y = graph.y[test_mask]
+        loss = self.default_loss_fn(pred, y)
+
+        self.note("test_loss", loss.item())
+        self.note("test_eval_index", pre_evaluation_index(pred, y))
+
+    def setup_optimizer(self):
+        cfg = self.optimizer_config
+        return torch.optim.Adam(self.model.parameters(), lr=cfg["lr"], weight_decay=cfg["weight_decay"])
diff --git a/cogdl/wrappers/model_wrapper/node_classification/sage_mw.py b/cogdl/wrappers/model_wrapper/node_classification/sage_mw.py
new file mode 100644
index 00000000..560c8177
--- /dev/null
+++ b/cogdl/wrappers/model_wrapper/node_classification/sage_mw.py
@@ -0,0 +1,46 @@
+import torch
+
+from .. import ModelWrapper, register_model_wrapper
+from cogdl.wrappers.tools.wrapper_utils import pre_evaluation_index
+
+
+@register_model_wrapper("sagn_mw")
+class SagnModelWrapper(ModelWrapper):
+    def __init__(self, model, optimizer_cfg):
+        super(SagnModelWrapper, self).__init__()
+        self.model = model
+
+    def train_step(self, batch):
+        batch_x, batch_y_emb, y = batch
+        pred = self.model(batch_x, batch_y_emb)
+        loss = self.default_loss_fn(pred, y[batch])
+        return loss
+
+    def val_step(self, batch):
+        batch_x, batch_y_emb, y = batch
+        pred = self.model(batch_x, batch_y_emb)
+
+        metric = self.evaluate(pred, y, metric="auto")
+
+        self.note("val_loss", self.default_loss_fn(pred, y))
+        self.note("val_metric", metric)
+
+    def test_step(self, batch):
+        pass
+
+    def post_stage(self, stage, data_w):
+        device = next(self.model.parameters()).device
+
+        self.model.eval()
+        preds = []
+
+        eval_loader = data_w.post_stage_wrapper()
+        with torch.no_grad():
+            for batch in eval_loader:
+                batch_x, batch_y_emb = data_w.post_stage_transform(batch)
+                batch_x = batch_x.to(device)
+                batch_y_emb = batch_y_emb.to(device) if batch_y_emb is not None else batch_y_emb
+                pred = self.model(batch_x, batch_y_emb)
+                preds.append(pred.to(self.data_device))
+        probs = torch.cat(preds, dim=0)
+        return probs
diff --git a/cogdl/wrappers/model_wrapper/node_classification/self_auxiliary_mw.py b/cogdl/wrappers/model_wrapper/node_classification/self_auxiliary_mw.py
new file mode 100644
index 00000000..42c37493
--- /dev/null
+++ b/cogdl/wrappers/model_wrapper/node_classification/self_auxiliary_mw.py
@@ -0,0 +1,85 @@
+import torch
+from .. import ModelWrapper, register_model_wrapper
+from cogdl.models.nn.self_auxiliary_task import (
+    EdgeMask,
+    PairwiseDistance,
+    PairwiseAttrSim,
+    AttributeMask,
+    Distance2Clusters,
+)
+from cogdl.wrappers.tools.wrapper_utils import evaluate_node_embeddings_using_logreg
+
+
+@register_model_wrapper("self_auxiliary_mw")
+class SelfAuxiliaryTask(ModelWrapper):
+    @staticmethod
+    def add_args(parser):
+        # fmt: off
+        parser.add_argument("--auxiliary-task", type=str, default="edge_mask",
+                            help="Option: edge_mask, attribute_mask, distance2clusters,"
+                                 " pairwise_distance, pairwise_attr_sim")
+        parser.add_argument("--dropedge-rate", type=float, default=0.0)
+        parser.add_argument("--mask-ratio", type=float, default=0.1)
+        # fmt: on
+
+    def __init__(self, model, optimizer_cfg, auxiliary_task, dropedge_rate, mask_ratio):
+        super().__init__()
+        self.auxiliary_task = auxiliary_task
+        self.optimizer_cfg = optimizer_cfg
+        self.hidden_size = optimizer_cfg["hidden_size"]
+        self.dropedge_rate = dropedge_rate
+        self.mask_ratio = mask_ratio
+        self.model = model
+
+        self.agent = None
+
+    def train_step(self, subgraph):
+        graph = subgraph
+        with graph.local_graph():
+            graph = self.agent.transform_data(graph)
+            pred = self.model(graph)
+        sup_loss = self.default_loss_fn(pred, graph.y)
+        ssl_loss = self.agent.make_loss(pred)
+        return sup_loss + ssl_loss
+
+    # def evaluate(self, dataset):
+    #     graph = dataset.data.to(self.device)
+
+    def test_step(self, graph):
+        with torch.no_grad():
+            pred = self.model(graph)
+        y = graph.y
+        result = evaluate_node_embeddings_using_logreg(pred, y, graph.train_mask, graph.test_mask)
+        self.note("test_acc", result)
+
+    def generate_virtual_labels(self, data):
+        if self.auxiliary_task == "edge_mask":
+            self.agent = EdgeMask(self.hidden_size, self.mask_ratio, self.device)
+        elif self.auxiliary_task == "attribute_mask":
+            self.agent = AttributeMask(data, self.hidden_size, data.train_mask, self.mask_ratio, self.device)
+        elif self.auxiliary_task == "pairwise_distance":
+            self.agent = PairwiseDistance(
+                self.hidden_size,
+                [(1, 2), (2, 3), (3, 5)],
+                self.sampling,
+                self.dropedge_rate,
+                256,
+                self.device,
+            )
+        elif self.auxiliary_task == "distance2clusters":
+            self.agent = Distance2Clusters(self.hidden_size, 30, self.device)
+        elif self.auxiliary_task == "pairwise_attr_sim":
+            self.agent = PairwiseAttrSim(self.hidden_size, 5, self.device)
+        else:
+            raise Exception(
+                "auxiliary task must be edge_mask, attribute_mask, pairwise_distance, distance2clusters,"
+                "or pairwise-attr-sim"
+            )
+
+    def self_supervised_loss(self, data):
+        embed = self.gcn.embed(data)
+        return self.agent.make_loss(embed)
+
+    def setup_optimizer(self):
+        lr, wd = self.optimizer_cfg["lr"], self.optimizer_cfg["weight_decay"]
+        return torch.optim.Adam(self.parameters(), lr=lr, weight_decay=wd)
diff --git a/cogdl/wrappers/model_wrapper/pretraining/__init__.py b/cogdl/wrappers/model_wrapper/pretraining/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/cogdl/wrappers/model_wrapper/pretraining/gcc_mw.py b/cogdl/wrappers/model_wrapper/pretraining/gcc_mw.py
new file mode 100644
index 00000000..3649beff
--- /dev/null
+++ b/cogdl/wrappers/model_wrapper/pretraining/gcc_mw.py
@@ -0,0 +1,133 @@
+import copy
+
+import torch
+import torch.nn as nn
+
+from .. import ModelWrapper, register_model_wrapper
+from cogdl.wrappers.tools.memory_moco import MemoryMoCo, NCESoftmaxLoss, moment_update
+from cogdl.utils.optimizer import LinearOptimizer
+
+
+@register_model_wrapper("gcc_mw")
+class GCCModelWrapper(ModelWrapper):
+    @staticmethod
+    def add_args(parser):
+        # loss function
+        parser.add_argument("--nce-k", type=int, default=32)
+        parser.add_argument("--nce-t", type=float, default=0.07)
+        parser.add_argument("--finetune", action="store_true")
+        parser.add_argument("--momentum", type=float, default=0.96)
+
+        # specify folder
+        parser.add_argument("--model-path", type=str, default=None, help="path to save model")
+
+    def __init__(
+        self,
+        model,
+        optimizer_cfg,
+        nce_k,
+        nce_t,
+        momentum,
+        output_size,
+        finetune=False,
+        num_classes=1,
+        model_path="gcc_pretrain.pt",
+    ):
+        super(GCCModelWrapper, self).__init__()
+        self.model = model
+        self.model_ema = copy.deepcopy(self.model)
+        for p in self.model_ema.parameters():
+            p.detach_()
+
+        self.optimizer_cfg = optimizer_cfg
+        self.output_size = output_size
+        self.momentum = momentum
+
+        self.contrast = MemoryMoCo(self.output_size, num_classes, nce_k, nce_t, use_softmax=True)
+        self.criterion = nn.CrossEntropyLoss() if finetune else NCESoftmaxLoss()
+
+        self.finetune = finetune
+        self.model_path = model_path
+        if finetune:
+            self.linear = nn.Linear(self.output_size, num_classes)
+        else:
+            self.register_buffer("linear", None)
+
+    def train_step(self, batch):
+        if self.finetune:
+            return self.train_step_finetune(batch)
+        else:
+            return self.train_step_pretraining(batch)
+
+    def train_step_pretraining(self, batch):
+        graph_q, graph_k = batch
+
+        # ===================Moco forward=====================
+        feat_q = self.model(graph_q)
+        with torch.no_grad():
+            feat_k = self.model_ema(graph_k)
+
+        out = self.contrast(feat_q, feat_k)
+
+        assert feat_q.shape == (graph_q.batch_size, self.output_size)
+        moment_update(self.model, self.model_ema, self.momentum)
+
+        loss = self.criterion(
+            out,
+        )
+        return loss
+
+    def train_step_finetune(self, batch):
+        graph = batch
+        y = graph.y
+        hidden = self.model(graph)
+        pred = self.linear(hidden)
+        loss = self.default_loss_fn(pred, y)
+        return loss
+
+    def setup_optimizer(self):
+        cfg = self.optimizer_cfg
+        lr = cfg["lr"]
+        weight_decay = cfg["weight_decay"]
+        warm_steps = cfg["n_warmup_steps"]
+        max_epoch = cfg["max_epoch"]
+        batch_size = cfg["batch_size"]
+        if self.finetune:
+            optimizer = torch.optim.Adam(
+                [{"params": self.model.parameters()}, {"params": self.linear.parameters()}],
+                lr=lr,
+                weight_decay=weight_decay,
+            )
+        else:
+            optimizer = torch.optim.Adam(self.model.parameters(), lr=lr, weight_decay=weight_decay)
+        optimizer = LinearOptimizer(optimizer, warm_steps, max_epoch * batch_size, init_lr=lr)
+        return optimizer
+
+    def save_checkpoint(self, path):
+        state = {
+            "model": self.model.state_dict(),
+            "contrast": self.contrast.state_dict(),
+            "model_ema": self.model_ema.state_dict(),
+        }
+        torch.save(state, path)
+
+    def load_checkpoint(self, path):
+        state = torch.load(path)
+        self.model.load_state_dict(state["model"])
+        self.model_ema.load_state_dict(state["model_ema"])
+        self.contrast.load_state_dict(state["contrast"])
+
+    def pre_stage(self, stage, data_w):
+        if self.finetune:
+            self.load_checkpoint(self.model_path)
+            self.model.apply(clear_bn)
+
+    def post_stage(self, stage, data_w):
+        if not self.finetune:
+            self.save_checkpoint(self.model_path)
+
+
+def clear_bn(m):
+    classname = m.__class__.__name__
+    if classname.find("BatchNorm") != -1:
+        m.reset_running_stats()
diff --git a/cogdl/wrappers/tools/__init__.py b/cogdl/wrappers/tools/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/cogdl/wrappers/tools/memory_moco.py b/cogdl/wrappers/tools/memory_moco.py
new file mode 100644
index 00000000..67e2263f
--- /dev/null
+++ b/cogdl/wrappers/tools/memory_moco.py
@@ -0,0 +1,83 @@
+import math
+
+import torch
+import torch.nn as nn
+
+
+class MemoryMoCo(nn.Module):
+    """Fixed-size queue with momentum encoder"""
+
+    def __init__(self, inputSize, outputSize, K, T=0.07, use_softmax=False):
+        super(MemoryMoCo, self).__init__()
+        self.outputSize = outputSize  # None
+        self.inputSize = inputSize
+        self.queueSize = K
+        self.T = T
+        self.index = 0
+        self.use_softmax = use_softmax
+
+        self.register_buffer("params", torch.tensor([-1]))
+        stdv = 1.0 / math.sqrt(inputSize / 3)
+        self.register_buffer("memory", torch.rand(self.queueSize, inputSize).mul_(2 * stdv).add_(-stdv))
+        print("using queue shape: ({},{})".format(self.queueSize, inputSize))
+
+    def forward(self, q, k):
+        batchSize = q.shape[0]
+        k = k.detach()
+
+        Z = self.params[0].item()
+
+        # pos logit
+        l_pos = torch.bmm(q.view(batchSize, 1, -1), k.view(batchSize, -1, 1))
+        l_pos = l_pos.view(batchSize, 1)
+        # neg logit
+        queue = self.memory.clone()
+        l_neg = torch.mm(queue.detach(), q.transpose(1, 0))
+        l_neg = l_neg.transpose(0, 1)
+
+        out = torch.cat((l_pos, l_neg), dim=1)
+
+        if self.use_softmax:
+            out = torch.div(out, self.T)
+            out = out.squeeze().contiguous()
+        else:
+            out = torch.exp(torch.div(out, self.T))
+            if Z < 0:
+                self.params[0] = out.mean() * self.outputSize
+                Z = self.params[0].clone().detach().item()
+                print("normalization constant Z is set to {:.1f}".format(Z))
+            # compute the out
+            out = torch.div(out, Z).squeeze().contiguous()
+
+        # # update memory
+        with torch.no_grad():
+            out_ids = torch.arange(batchSize).cuda()
+            out_ids += self.index
+            out_ids = torch.fmod(out_ids, self.queueSize)
+            out_ids = out_ids.long()
+            self.memory.index_copy_(0, out_ids, k)
+            self.index = (self.index + batchSize) % self.queueSize
+
+        return out
+
+
+class NCESoftmaxLoss(nn.Module):
+    """Softmax cross-entropy loss (a.k.a., info-NCE loss in CPC paper)"""
+
+    def __init__(self):
+        super(NCESoftmaxLoss, self).__init__()
+        self.criterion = nn.CrossEntropyLoss()
+
+    def forward(self, x):
+        bsz = x.shape[0]
+        x = x.squeeze()
+        label = torch.zeros([bsz]).cuda().long()
+        loss = self.criterion(x, label)
+        return loss
+
+
+def moment_update(model, model_ema, m):
+    """ model_ema = m * model_ema + (1 - m) model """
+    for p1, p2 in zip(model.parameters(), model_ema.parameters()):
+        # p2.data.mul_(m).add_(1 - m, p1.detach().data)
+        p2.data.mul_(m).add_(p1.detach().data * (1 - m))
diff --git a/cogdl/wrappers/tools/wrapper_utils.py b/cogdl/wrappers/tools/wrapper_utils.py
new file mode 100644
index 00000000..3c2a6eeb
--- /dev/null
+++ b/cogdl/wrappers/tools/wrapper_utils.py
@@ -0,0 +1,292 @@
+from typing import Dict
+
+import random
+import numpy as np
+import scipy.sparse as sp
+from collections import defaultdict
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from sklearn.utils import shuffle as skshuffle
+from sklearn.linear_model import LogisticRegression
+from sklearn.multiclass import OneVsRestClassifier
+from sklearn.model_selection import GridSearchCV, KFold
+from sklearn.svm import SVC
+from sklearn.cluster import KMeans, SpectralClustering
+from sklearn.metrics.cluster import normalized_mutual_info_score
+from sklearn.metrics import f1_score
+from scipy.optimize import linear_sum_assignment
+
+from cogdl.utils import accuracy, multilabel_f1
+
+
+def pre_evaluation_index(y_pred, y_true, sigmoid=False):
+    """
+    Pre-calculating diffusion matrix for mini-batch evaluation
+    Return:
+        torch.Tensor((tp, all)) for multi-class classification
+        torch.Tensor((tp, fp, fn)) for multi-label classification
+    """
+    if len(y_true.shape) == 1:
+        pred = (y_pred.argmax(1) == y_true).int()
+        tp = pred.sum()
+        fnp = pred.shape[0] - tp
+        return torch.tensor((tp, fnp)).float()
+    else:
+        if sigmoid:
+            border = 0.5
+        else:
+            border = 0
+        y_pred[y_pred >= border] = 1
+        y_pred[y_pred < border] = 0
+        tp = (y_pred * y_true).sum().to(torch.float32)
+        # tn = ((1 - y_true) * (1 - y_pred)).sum().to(torch.float32)
+        fp = ((1 - y_true) * y_pred).sum().to(torch.float32)
+        fn = (y_true * (1 - y_pred)).sum().to(torch.float32)
+        return torch.tensor((tp, fp, fn))
+
+
+def merge_batch_indexes(values: list, method="mean"):
+    # if key.endswith("loss"):
+    #     result = sum(values)
+    #     if torch.is_tensor(result):
+    #         result = result.item()
+    #     result = result / len(values)
+    # elif key.endswith("eval_index"):
+    #     if len(values) > 1:
+    #         val = torch.stack(values)
+    #         val = val.sum(0)
+    #     else:
+    #         val = values[0]
+    #     fp = val[0]
+    #     all_ = val.sum()
+    #
+    #     prefix = key[: key.find("eval_index")]
+    #     if val.shape[0] == 2:
+    #         _key = prefix + "acc"
+    #     else:
+    #         _key = prefix + "f1"
+    #     result = (fp / all_).item()
+    #     out_key = _key
+
+    if isinstance(values[0], dict) or isinstance(values[0], tuple):
+        return values
+    elif method == "mean":
+        return sum(values) / len(values)
+    elif method == "sum":
+        return sum(values)
+    else:
+        return sum(values)
+
+
+def node_degree_as_feature(data):
+    r"""
+    Set each node feature as one-hot encoding of degree
+    :param data: a list of class Data
+    :return: a list of class Data
+    """
+    max_degree = 0
+    degrees = []
+    device = data[0].edge_index[0].device
+
+    for graph in data:
+        deg = graph.degrees()
+        degrees.append(deg)
+        max_degree = max(deg.max().item(), max_degree)
+
+    max_degree = int(max_degree) + 1
+    for i in range(len(data)):
+        one_hot = F.one_hot(degrees[i], max_degree)
+        data[i].x = one_hot.to(device)
+    return data
+
+
+def split_dataset(ndata, train_ratio, test_ratio):
+
+    train_size = int(ndata * train_ratio)
+    test_size = int(ndata * test_ratio)
+    index = np.arange(ndata)
+    random.shuffle(index)
+
+    train_index = index[:train_size]
+    test_index = index[-test_size:]
+    if train_ratio + test_ratio == 1:
+        val_index = None
+    else:
+        val_index = index[train_size:-test_size]
+    return train_index, val_index, test_index
+
+
+def evaluate_node_embeddings_using_logreg(data, labels, train_idx, test_idx, run=20):
+    result = LogRegTrainer().train(data, labels, train_idx, test_idx, run=run)
+    return result
+
+
+class LogReg(nn.Module):
+    def __init__(self, ft_in, nb_classes):
+        super(LogReg, self).__init__()
+        self.fc = nn.Linear(ft_in, nb_classes)
+
+        for m in self.modules():
+            self.weights_init(m)
+
+    def weights_init(self, m):
+        if isinstance(m, nn.Linear):
+            torch.nn.init.xavier_uniform_(m.weight.data)
+            if m.bias is not None:
+                m.bias.data.fill_(0.0)
+
+    def forward(self, seq):
+        ret = self.fc(seq)
+        return ret
+
+
+class LogRegTrainer(object):
+    def train(self, data, labels, idx_train, idx_test, loss_fn=None, evaluator=None, run=20):
+        device = data.device
+        nhid = data.shape[-1]
+        labels = labels.to(device)
+
+        train_embs = data[idx_train]
+        test_embs = data[idx_test]
+
+        train_lbls = labels[idx_train]
+        test_lbls = labels[idx_test]
+        tot = 0
+
+        num_classes = int(labels.max()) + 1
+
+        if loss_fn is None:
+            loss_fn = nn.CrossEntropyLoss() if len(labels.shape) == 1 else nn.BCEWithLogitsLoss()
+
+        if evaluator is None:
+            evaluator = accuracy if len(labels.shape) == 1 else multilabel_f1
+
+        for _ in range(run):
+            log = LogReg(nhid, num_classes).to(device)
+            optimizer = torch.optim.Adam(log.parameters(), lr=0.01, weight_decay=0.0)
+            log.to(device)
+
+            for _ in range(100):
+                log.train()
+                optimizer.zero_grad()
+
+                logits = log(train_embs)
+                loss = loss_fn(logits, train_lbls)
+
+                loss.backward()
+                optimizer.step()
+
+            log.eval()
+            with torch.no_grad():
+                logits = log(test_embs)
+            metric = evaluator(logits, test_lbls)
+
+            tot += metric
+        return tot / run
+
+
+def evaluate_node_embeddings_using_liblinear(features_matrix, label_matrix, num_shuffle, training_percents):
+    if len(label_matrix.shape) > 1:
+        labeled_nodes = np.nonzero(np.sum(label_matrix, axis=1) > 0)[0]
+        features_matrix = features_matrix[labeled_nodes]
+        label_matrix = label_matrix[labeled_nodes]
+
+    # shuffle, to create train/test groups
+    shuffles = []
+    for _ in range(num_shuffle):
+        shuffles.append(skshuffle(features_matrix, label_matrix))
+
+    # score each train/test group
+    all_results = defaultdict(list)
+
+    for train_percent in training_percents:
+        for shuf in shuffles:
+            X, y = shuf
+
+            training_size = int(train_percent * len(features_matrix))
+            X_train = X[:training_size, :]
+            y_train = y[:training_size, :]
+
+            X_test = X[training_size:, :]
+            y_test = y[training_size:, :]
+
+            clf = TopKRanker(LogisticRegression(solver="liblinear"))
+            clf.fit(X_train, y_train)
+
+            # find out how many labels should be predicted
+            top_k_list = y_test.sum(axis=1).astype(np.int).tolist()
+            preds = clf.predict(X_test, top_k_list)
+            result = f1_score(y_test, preds, average="micro")
+            all_results[train_percent].append(result)
+
+    return dict(
+        (f"Micro-F1 {train_percent}", np.mean(all_results[train_percent]))
+        for train_percent in sorted(all_results.keys())
+    )
+
+
+class TopKRanker(OneVsRestClassifier):
+    def predict(self, X, top_k_list):
+        assert X.shape[0] == len(top_k_list)
+        probs = np.asarray(super(TopKRanker, self).predict_proba(X))
+        all_labels = sp.lil_matrix(probs.shape)
+
+        for i, k in enumerate(top_k_list):
+            probs_ = probs[i, :]
+            labels = self.classes_[probs_.argsort()[-k:]].tolist()
+            for label in labels:
+                all_labels[i, label] = 1
+        return all_labels
+
+
+def evaluate_graph_embeddings_using_svm(embeddings, labels):
+    result = []
+    kf = KFold(n_splits=10)
+    kf.get_n_splits(X=embeddings, y=labels)
+    for train_index, test_index in kf.split(embeddings):
+        x_train = embeddings[train_index]
+        x_test = embeddings[test_index]
+        y_train = labels[train_index]
+        y_test = labels[test_index]
+        params = {"C": [1e-2, 1e-1, 1]}
+        svc = SVC()
+        clf = GridSearchCV(svc, params)
+        clf.fit(x_train, y_train)
+
+        preds = clf.predict(x_test)
+        f1 = f1_score(y_test, preds, average="micro")
+        result.append(f1)
+    test_f1 = np.mean(result)
+    test_std = np.std(result)
+
+    return dict(acc=test_f1, std=test_std)
+
+
+def evaluate_clustering(features_matrix, labels, cluster_method, num_clusters, num_nodes, full=True):
+    print("Clustering...")
+    if cluster_method == "kmeans":
+        kmeans = KMeans(n_clusters=num_clusters, random_state=0).fit(features_matrix)
+        clusters = kmeans.labels_
+    else:
+        clustering = SpectralClustering(n_clusters=num_clusters, assign_labels="discretize", random_state=0).fit(
+            features_matrix
+        )
+        clusters = clustering.labels_
+
+    print("Evaluating...")
+    truth = labels.cpu().numpy()
+    if full:
+        mat = np.zeros([num_clusters, num_clusters])
+        for i in range(num_nodes):
+            mat[clusters[i]][truth[i]] -= 1
+        _, row_idx = linear_sum_assignment(mat)
+        acc = -mat[_, row_idx].sum() / num_nodes
+        for i in range(num_nodes):
+            clusters[i] = row_idx[clusters[i]]
+        macro_f1 = f1_score(truth, clusters, average="macro")
+        return dict(Accuracy=acc, NMI=normalized_mutual_info_score(clusters, truth), Macro_F1=macro_f1)
+    else:
+        return dict(NMI=normalized_mutual_info_score(clusters, truth))
diff --git a/examples/custom_dataset.py b/examples/custom_dataset.py
index 7f536082..e3c2b3bd 100644
--- a/examples/custom_dataset.py
+++ b/examples/custom_dataset.py
@@ -1,6 +1,6 @@
 import torch
 
-from cogdl import experiment
+from cogdl.runner.experiments import experiment
 from cogdl.data import Graph
 from cogdl.datasets import NodeDataset, register_dataset
 
@@ -8,7 +8,6 @@
 @register_dataset("mydataset")
 class MyNodeClassificationDataset(NodeDataset):
     def __init__(self):
-        self.path = "mydata.pt"
         super(MyNodeClassificationDataset, self).__init__(self.path)
 
     def process(self):
@@ -29,12 +28,11 @@ def process(self):
         test_mask = torch.zeros(num_nodes).bool()
         test_mask[int(0.7 * num_nodes) :] = True
         data = Graph(x=x, edge_index=edge_index, y=y, train_mask=train_mask, val_mask=val_mask, test_mask=test_mask)
-        torch.save(data, self.path)
         return data
 
 
 if __name__ == "__main__":
     # Run with self-loaded dataset
-    experiment(task="node_classification", dataset="mydataset", model="gcn")
-    # Run with given datapaath
-    experiment(task="node_classification", dataset="./mydata.pt", model="gcn")
+    experiment(dw="node_classification_dw", mw="node_classification_mw", dataset="mydataset", model="gcn")
+    # Run with given datapath
+    experiment(dw="node_classification_dw", mw="node_classification_mw", dataset="./mydata.pt", model="gcn")
diff --git a/examples/quick_start.py b/examples/quick_start.py
index ceb6cb8a..275dd4db 100644
--- a/examples/quick_start.py
+++ b/examples/quick_start.py
@@ -1,13 +1,13 @@
 from cogdl import experiment
 
-# basic usage
-experiment(task="node_classification", dataset="cora", model="gcn")
+# # basic usage
+# experiment(dataset="cora", model="gcn")
 
-# set other hyper-parameters
-experiment(task="node_classification", dataset="cora", model="gcn", hidden_size=32, max_epoch=200)
+# # set other hyper-parameters
+# experiment(dataset="cora", model="gcn", hidden_size=32, max_epoch=200)
 
-# run over multiple models on different seeds
-experiment(task="node_classification", dataset="cora", model=["gcn", "gat"], seed=[1, 2])
+# # run over multiple models on different seeds
+# experiment(dataset="cora", model=["gcn", "gat"], seed=[1, 2])
 
 
 # automl usage
@@ -19,4 +19,4 @@ def func_search(trial):
     }
 
 
-experiment(task="node_classification", dataset="cora", model="gcn", seed=[1, 2], func_search=func_search)
+experiment(dataset="cora", model="gcn", seed=[1, 2], func_search=func_search)
diff --git a/scripts/main.py b/scripts/main.py
new file mode 100644
index 00000000..09646e17
--- /dev/null
+++ b/scripts/main.py
@@ -0,0 +1,11 @@
+import cogdl.runner.options as options
+from cogdl.runner.experiments import experiment
+
+
+if __name__ == "__main__":
+    parser = options.get_training_parser()
+    args, _ = parser.parse_known_args()
+    args = options.parse_args_and_arch(parser, args)
+    print(args)
+
+    experiment(dataset=args.dataset, model=args.model, args=args)


Test patch:
diff --git a/tests/models/ssl/test_generative_models.py b/tests/models/ssl/test_generative_models.py
index 3a70c16d..d07d0885 100644
--- a/tests/models/ssl/test_generative_models.py
+++ b/tests/models/ssl/test_generative_models.py
@@ -12,7 +12,7 @@ def get_default_args():
     cuda_available = torch.cuda.is_available()
     args = {
         "dataset": "cora",
-        "trainer": "self_supervised_joint",
+        "controller": "self_supervised_joint",
         "model": "self_auxiliary_task",
         "hidden_size": 64,
         "label_mask": 0,