Repository: THUDM/CogDL (language: python)

Issue:
random.seed
I want to use the function _display_data.py_, I get the error:
```
Traceback (most recent call last):
  File "display_data.py", line 75, in <module>
    random.seed(args.seed)
  File "/Users/wangzhikai/.conda/envs/cogdl/lib/python3.7/random.py", line 126, in seed
    super().seed(a)
TypeError: unhashable type: 'list'
```



Discussion / hints:
Hi @Cloudcatcher888 ,

Thanks for using our cogdl and pointing this out. For _display_data.py_, it only supports passing one integer to the _seed_. It seems that the value of _seed_ is a list (default: [1]), which causes this error. We will deal with the inconsistent bug of the code immediately. Please accept my apologies for this mistake.

Fix patch:
diff --git a/.travis.yml b/.travis.yml
index cf18870a..481a4504 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -23,7 +23,6 @@ before_script:
 script:
   - export DGLBACKEND=pytorch
   - export METIS_DLL=~/.local/lib/libmetis.so
-  - ls ~/.local/lib && echo $METIS_DLL
   - coverage run --source=cogdl -m pytest tests
   - cd docs && make clean && make html && cd ..
 
diff --git a/cogdl/data/data.py b/cogdl/data/data.py
index 88a6e891..077fc72b 100644
--- a/cogdl/data/data.py
+++ b/cogdl/data/data.py
@@ -218,6 +218,13 @@ def _build_adj_(self):
         self.__adj = sparse.csr_matrix(
             (edge_attr_np, (edge_index_np[0], edge_index_np[1])), shape=(num_nodes, num_nodes)
         )
+        return self.__adj
+
+    def _eliminate_adj_(self):
+        self.__adj = None
+        self.edge_row = None
+        self.edge_col = None
+        self.node_idx = None
 
     def subgraph(self, node_idx):
         """Return the induced node subgraph."""
@@ -240,7 +247,7 @@ def subgraph(self, node_idx):
             attrs["edge_attr"] = edge_attr
         return Data(**attrs)
 
-    def edge_subgraph(self, edge_idx):
+    def edge_subgraph(self, edge_idx, require_idx=False):
         """Return the induced edge subgraph."""
         if isinstance(edge_idx, torch.Tensor):
             edge_idx = edge_idx.cpu().numpy()
@@ -260,6 +267,8 @@ def func(x):
         attrs["edge_index"] = edge_index
         if edge_attr is not None:
             attrs["edge_attr"] = edge_attr
+        if require_idx:
+            return Data(**attrs), node_idx, edge_idx
         return Data(**attrs)
 
     def sample_adj(self, batch, size=-1, replace=True):
diff --git a/cogdl/data/sampler.py b/cogdl/data/sampler.py
index 46704555..1e9d04b7 100644
--- a/cogdl/data/sampler.py
+++ b/cogdl/data/sampler.py
@@ -7,7 +7,26 @@
 import torch.utils.data
 
 from cogdl.data import Data
-from cogdl.utils import remove_self_loops
+from cogdl.utils import remove_self_loops, row_normalization
+
+
+def normalize(adj):
+    D = adj.sum(1).flatten()
+    norm_diag = sp.dia_matrix((1 / D, 0), shape=adj.shape)
+    adj = norm_diag.dot(adj)
+    adj.sort_indices()
+    return adj
+
+
+def _coo_scipy2torch(adj):
+    """
+    convert a scipy sparse COO matrix to torch
+    """
+    values = adj.data
+    indices = np.vstack((adj.row, adj.col))
+    i = torch.LongTensor(indices)
+    v = torch.FloatTensor(values)
+    return torch.sparse.FloatTensor(i, v, torch.Size(adj.shape))
 
 
 class Sampler:
@@ -19,6 +38,7 @@ class Sampler:
 
     def __init__(self, data, args_params):
         self.data = data.clone()
+        self.full_graph = data.clone()
         self.num_nodes = self.data.x.size()[0]
         self.num_edges = (
             self.data.edge_index_train.size()[1]
@@ -46,23 +66,29 @@ class SAINTSampler(Sampler):
 
     def __init__(self, data, args_params):
         super().__init__(data, args_params)
+
+        self.gen_adj()
+
+        self.train_mask = self.data.train_mask.cpu().numpy()
+        self.node_train = np.arange(1, self.num_nodes + 1) * self.train_mask
+        self.node_train = self.node_train[self.node_train != 0] - 1
+
+        self.sample_coverage = args_params["sample_coverage"]
+        self.preprocess()
+
+    def gen_adj(self):
         edge_index = (
             self.data.edge_index_train.cpu().numpy()
             if hasattr(self.data, "edge_index_train")
             else self.data.edge_index.cpu().numpy()
         )
+
         self.adj = sp.coo_matrix(
             (np.ones(self.num_edges), (edge_index[0], edge_index[1])),
             shape=(self.num_nodes, self.num_nodes),
         ).tocsr()
-        self.train_mask = self.data.train_mask.cpu().numpy()
-        self.node_train = np.arange(1, self.num_nodes + 1) * self.train_mask
-        self.node_train = self.node_train[self.node_train != 0] - 1
-
-        self.sample_coverage = args_params["sample_coverage"]
-        self.estimate()
 
-    def estimate(self):
+    def preprocess(self):
         r"""
         estimation of loss / aggregation normalization factors.
         For some special sampler, no need to estimate norm factors, we can calculate
@@ -73,32 +99,31 @@ def estimate(self):
             2. update the counter for each node / edge in the training graph
             3. estimate norm factor alpha and lambda
         """
-        self.subgraphs_indptr = []
-        self.subgraphs_indices = []
-        self.subgraphs_data = []
-        self.subgraphs_nodes = []
-        self.subgraphs_edge_index = []
+        self.subgraph_data = []
+        self.subgraph_node_idx = []
+        self.subgraph_edge_idx = []
 
         self.norm_loss_train = np.zeros(self.num_nodes)
         self.norm_aggr_train = np.zeros(self.num_edges)
         self.norm_loss_test = np.ones(self.num_nodes) / self.num_nodes
         self.norm_loss_test = torch.from_numpy(self.norm_loss_test.astype(np.float32))
 
-        tot_sampled_nodes = 0
+        num_sampled_nodes = 0
         while True:
-            tot_sampled_nodes += self.gen_subgraph()
+            num_sampled_nodes += self.gen_subgraph()
             print(
                 "\rGenerating subgraphs %.2lf%%"
-                % (tot_sampled_nodes * 100 / self.data.num_nodes / self.sample_coverage),
+                % min(num_sampled_nodes * 100 / self.data.num_nodes / self.sample_coverage, 100),
                 end="",
                 flush=True,
             )
-            if tot_sampled_nodes > self.sample_coverage * self.num_nodes:
+            if num_sampled_nodes > self.sample_coverage * self.num_nodes:
                 break
-        num_subg = len(self.subgraphs_nodes)
+
+        num_subg = len(self.subgraph_data)
         for i in range(num_subg):
-            self.norm_aggr_train[self.subgraphs_edge_index[i]] += 1
-            self.norm_loss_train[self.subgraphs_nodes[i]] += 1
+            self.norm_aggr_train[self.subgraph_edge_idx[i]] += 1
+            self.norm_loss_train[self.subgraph_node_idx[i]] += 1
         for v in range(self.data.num_nodes):
             i_s = self.adj.indptr[v]
             i_e = self.adj.indptr[v + 1]
@@ -109,126 +134,84 @@ def estimate(self):
         self.norm_loss_train[self.node_train] = num_subg / self.norm_loss_train[self.node_train] / self.node_train.size
         self.norm_loss_train = torch.from_numpy(self.norm_loss_train.astype(np.float32))
 
-    def gen_subgraph(self):
-        _indptr, _indices, _data, _v, _edge_index = self.sample()
-        self.subgraphs_indptr.append(_indptr)
-        self.subgraphs_indices.append(_indices)
-        self.subgraphs_data.append(_data)
-        self.subgraphs_nodes.append(_v)
-        self.subgraphs_edge_index.append(_edge_index)
-        return len(_v)
-
-    def sample(self):
-        pass
-
-    def extract_subgraph(self, edge_idx, directed=True):
-        edge_idx = np.unique(edge_idx)
-        subg_edge = (
-            self.data.edge_index_train.t()[edge_idx].cpu().numpy()
-            if hasattr(self.data, "edge_index_train")
-            else self.data.edge_index.t()[edge_idx].cpu().numpy()
-        )
-        if not directed:
-            subg_edge = np.concatenate((subg_edge, subg_edge[:, [1, 0]]))
-        subg_edge = np.unique(subg_edge, axis=0)
-        # get nodes whose degree != 0
-        nodes = np.zeros(self.num_nodes, dtype=int)
-        for e in subg_edge:
-            nodes[e[0]] = 1
-            nodes[e[1]] = 1
-        if not directed:
-            edge_idx = []
-            for e in subg_edge:
-                for i in range(self.adj.indptr[e[0]], self.adj.indptr[e[0] + 1]):
-                    if self.adj.indices[i] == e[1]:
-                        edge_idx.append(i)
-            edge_idx = np.array(edge_idx, dtype=int)
-        node_idx = np.arange(1, self.num_nodes + 1) * np.array(nodes)
-        node_idx = node_idx[node_idx != 0] - 1
-
-        # mapping nodes to new indices
-        orig2subg = {n: i for i, n in enumerate(node_idx)}
-        indptr = np.zeros(node_idx.size + 1, dtype=int)
-        for ind, u in enumerate(subg_edge[:, 0]):
-            if ind + 1 > indptr[orig2subg[u] + 1]:
-                indptr[orig2subg[u] + 1] = ind + 1
-        for i in range(1, node_idx.size + 1):
-            if indptr[i] == 0:
-                indptr[i] = indptr[i - 1]
-        indices = subg_edge[:, 1]
-        for i in range(len(indices)):
-            indices[i] = orig2subg[indices[i]]
-        data = np.ones(indices.size)
-        assert indptr[-1] == indices.size == subg_edge.size // 2
-        return indptr, indices, data, node_idx, edge_idx
-
-    def get_subgraph(self, phase, require_norm=True):
+    def one_batch(self, phase, require_norm=True):
         r"""
         Generate one minibatch for model. In the 'train' mode, one minibatch corresponds
-        to one subgraph of the training graph. In the 'valid' or 'test' mode, one batch
+        to one subgraph of the training graph. In the 'val' or 'test' mode, one batch
         corresponds to the full graph (i.e., full-batch rather than minibatch evaluation
         for validation / test sets).
 
         Inputs:
-            mode                str, can be 'train', 'valid', 'test'
+            phase               str, can be 'train', 'val', 'test'
             require_norm        boolean
 
         Outputs:
             data                Data object, modeling the sampled subgraph
             data.norm_aggr      aggregation normalization
-            data.norm_loss      normalization normalization
+            data.norm_loss      loss normalization
         """
         if phase in ["val", "test"]:
-            node_subgraph = np.arange(self.data.num_nodes)
-            data = self.data.clone()
-            if require_norm:
-                data.norm_aggr = torch.ones(self.data.edge_index.size()[1])
-                data.norm_loss = self.norm_loss_test
+            data = self.full_graph.clone()
+            data.norm_loss = self.norm_loss_test
         else:
             while True:
-                if len(self.subgraphs_nodes) == 0:
+                if len(self.subgraph_data) == 0:
                     self.gen_subgraph()
 
-                node_subgraph = self.subgraphs_nodes.pop()
-                edge_subgraph = self.subgraphs_edge_index.pop()
-                flag = False
-                for idx in node_subgraph:
-                    if self.train_mask[idx]:
-                        flag = True
-                        break
-                if flag:
+                data = self.subgraph_data.pop()
+                node_idx = self.subgraph_node_idx.pop()
+                edge_idx = self.subgraph_edge_idx.pop()
+                if self.exists_train_nodes(node_idx):
                     break
-            num_nodes_subgraph = node_subgraph.size
-            adj = sp.csr_matrix(
-                (self.subgraphs_data.pop(), self.subgraphs_indices.pop(), self.subgraphs_indptr.pop()),
-                shape=(num_nodes_subgraph, num_nodes_subgraph),
-            )
 
             if require_norm:
-                adj.data[:] = self.norm_aggr_train[edge_subgraph][:]
-                # normalization
-                D = adj.sum(1).flatten()
-                norm_diag = sp.dia_matrix((1 / D, 0), shape=adj.shape)
-                adj = norm_diag.dot(adj)
-                adj.sort_indices()
-
-            adj = adj.tocoo()
-            data = Data(
-                self.data.x[node_subgraph],
-                torch.LongTensor(np.vstack((adj.row, adj.col))),
-                None if self.data.edge_attr is None else self.data.edge_attr[edge_subgraph],
-                self.data.y[node_subgraph],
-                None if self.data.pos is None else self.data.pos[node_subgraph],
-            )
+                data.norm_aggr = torch.FloatTensor(self.norm_aggr_train[edge_idx][:])
+                data.norm_loss = self.norm_loss_train[node_idx]
+
+            data.train_mask = self.data.train_mask[node_idx]
+            data.val_mask = self.data.val_mask[node_idx]
+            data.test_mask = self.data.test_mask[node_idx]
+
+        # adj = data._build_adj_()
+        # adj = normalize(adj).tocoo()
+        # data.adj = _coo_scipy2torch(adj)
+        edge_weight = row_normalization(data.x.shape[0], data.edge_index)
+        data.edge_weight = edge_weight
+        # data._eliminate_adj_()
+        return data
 
-            if require_norm:
-                data.norm_aggr = torch.FloatTensor(adj.data)
-                data.norm_loss = self.norm_loss_train[node_subgraph]
-            data.train_mask = self.data.train_mask[node_subgraph]
-            data.val_mask = self.data.val_mask[node_subgraph]
-            data.test_mask = self.data.test_mask[node_subgraph]
+    def exists_train_nodes(self, node_idx):
+        for idx in node_idx:
+            if self.train_mask[idx]:
+                return True
+        return False
 
-        return data
+    def node_induction(self, node_idx):
+        node_idx = np.unique(node_idx)
+        node_flags = np.zeros(self.num_nodes)
+        for u in node_idx:
+            node_flags[u] = 1
+        edge_idx = []
+        for u in node_idx:
+            for e in range(self.adj.indptr[u], self.adj.indptr[u + 1]):
+                v = self.adj.indices[e]
+                if node_flags[v]:
+                    edge_idx.append(e)
+        edge_idx = np.array(edge_idx)
+        return self.data.subgraph(node_idx), node_idx, edge_idx
+
+    def edge_induction(self, edge_idx):
+        return self.data.edge_subgraph(edge_idx, require_idx=True)
+
+    def gen_subgraph(self):
+        _data, _node_idx, _edge_idx = self.sample()
+        self.subgraph_data.append(_data)
+        self.subgraph_node_idx.append(_node_idx)
+        self.subgraph_edge_idx.append(_edge_idx)
+        return len(_node_idx)
+
+    def sample(self):
+        pass
 
 
 class NodeSampler(SAINTSampler):
@@ -245,26 +228,7 @@ def __init__(self, data, args_params):
 
     def sample(self):
         node_idx = np.random.choice(np.arange(self.num_nodes), self.node_num_subgraph)
-        node_idx = np.unique(node_idx)
-        node_idx.sort()
-        orig2subg = {n: i for i, n in enumerate(node_idx)}
-        indptr = np.zeros(node_idx.size + 1)
-        indices = []
-        subg_edge_index = []
-        for nid in node_idx:
-            idx_s, idx_e = self.adj.indptr[nid], self.adj.indptr[nid + 1]
-            neighs = self.adj.indices[idx_s:idx_e]
-            for i_n, n in enumerate(neighs):
-                if n in orig2subg:
-                    indices.append(orig2subg[n])
-                    indptr[orig2subg[nid] + 1] += 1
-                    subg_edge_index.append(idx_s + i_n)
-        indptr = indptr.cumsum().astype(np.int64)
-        indices = np.array(indices)
-        subg_edge_index = np.array(subg_edge_index)
-        data = np.ones(indices.size)
-        assert indptr[-1] == indices.size == subg_edge_index.size
-        return indptr, indices, data, node_idx, subg_edge_index
+        return self.node_induction(node_idx)
 
 
 class EdgeSampler(SAINTSampler):
@@ -281,12 +245,17 @@ def __init__(self, data, args_params):
 
     def sample(self):
         edge_idx = np.random.choice(np.arange(self.num_edges), self.edge_num_subgraph)
-        return self.extract_subgraph(edge_idx)
+        return self.edge_induction(edge_idx)
 
 
 class RWSampler(SAINTSampler):
     r"""
-    randomly select a node, perform a random walk starting from the node and add the walk path to the subgraph
+    The sampler performs unbiased random walk, by following the steps:
+     1. Randomly pick `size_root` number of root nodes from all training nodes;
+     2. Perform length `size_depth` random walk from the roots. The current node
+            expands the next hop by selecting one of the neighbors uniformly
+            at random;
+     3. Generate node-induced subgraph from the nodes touched by the random walk.
     Args:
         sample_coverage (integer):  number of sampled nodes during estimation / number of nodes in graph
         num_walks (integer): number of walks
@@ -299,29 +268,24 @@ def __init__(self, data, args_params):
         super().__init__(data, args_params)
 
     def sample(self):
-        edge_idx = []
+        node_idx = []
         for walk in range(self.num_walks):
-            u = np.random.randint(self.num_nodes)
+            u = np.random.choice(self.node_train)
+            node_idx.append(u)
             for step in range(self.walk_length):
                 idx_s = self.adj.indptr[u]
                 idx_e = self.adj.indptr[u + 1]
                 if idx_s >= idx_e:
                     break
                 e = np.random.randint(idx_s, idx_e)
-                edge_idx.append(e)
                 u = self.adj.indices[e]
+                node_idx.append(u)
 
-        return self.extract_subgraph(np.array(edge_idx))
+        return self.node_induction(node_idx)
 
 
 class MRWSampler(SAINTSampler):
-    r"""
-    multidimentional random walk, similar to https://arxiv.org/abs/1002.1751
-    Args:
-        sample_coverage (integer):  number of sampled nodes during estimation / number of nodes in graph
-        size_frontier (integer): number of frontiers
-        size_subgraph (integer): number of edges in subgraph
-    """
+    r"""multidimentional random walk, similar to https://arxiv.org/abs/1002.1751"""
 
     def __init__(self, data, args_params):
         self.size_frontier = args_params["size_frontier"]
@@ -341,6 +305,8 @@ def sample(self):
                 id += 1
             nid = frontier[id]
             idx_s, idx_e = self.adj.indptr[nid], self.adj.indptr[nid + 1]
+            if idx_s >= idx_e:
+                continue
             e = np.random.randint(idx_s, idx_e)
             edge_idx.append(e)
             v = self.adj.indices[e]
@@ -349,7 +315,7 @@ def sample(self):
             deg[id] = self.adj.indptr[v + 1] - self.adj.indptr[v]
             deg_sum += deg[id]
 
-        return self.extract_subgraph(np.array(edge_idx))
+        return self.edge_induction(np.array(edge_idx))
 
 
 class LayerSampler(Sampler):
@@ -455,124 +421,3 @@ def batcher(self, batch):
         nodes = np.concatenate([self.clusters[i] for i in batch])
         subgraph = self.data.subgraph(nodes)
         return subgraph
-
-
-"""class FastGCNSampler(LayerSampler):
-    def __init__(self, data, params_args):
-        super().__init__(data, params_args)
-
-    def generate_adj(self, sample1, sample2):
-        edgelist = []
-        mapping = {}
-        for i in range(len(sample1)):
-            mapping[sample1[i]] = i
-
-        for i in range(len(sample2)):
-            nodes = self.adj[sample2[i]]
-            for node in nodes:
-                if node in mapping:
-                    edgelist.append([mapping[node], i])
-        edgetensor = torch.LongTensor(edgelist)
-        valuetensor = torch.ones(edgetensor.shape[0]).float()
-        t = torch.sparse_coo_tensor(
-            edgetensor.t(), valuetensor, (len(sample1), len(sample2))
-        )
-        return t
-
-    def sample_one_layer(self, sampled, sample_size):
-        total = []
-        for node in sampled:
-            total.extend(self.adj[node])
-        total = list(set(total))
-        if sample_size < len(total):
-            total = random.sample(total, sample_size)
-        return total
-
-    def sample(self, x, v, num_layers):
-        all_support = [[] for _ in range(num_layers)]
-        sampled = v.detach().cpu().numpy()
-        for i in range(num_layers - 1, -1, -1):
-            cur_sampled = self.sample_one_layer(sampled, self.sample_size[i])
-            all_support[i] = self.generate_adj(sampled, cur_sampled).to(x.device)
-            sampled = cur_sampled
-
-        return x[torch.LongTensor(sampled).to(x.device)], all_support, 0
-
-class ASGCNSampler(LayerSampler):
-    def __init__(self, data, params_args):
-        super().__init__(data, params_args)
-
-    def set_w(w_s0, w_s1):
-        self.w_s0 = w_s0
-        self.w_s1 = w_s1
-
-    def set_adj(self, edge_index, num_nodes):
-        self.sparse_adj = sparse.coo_matrix(
-            (np.ones(edge_index.shape[1]), (edge_index[0], edge_index[1])),
-            shape=(num_nodes, num_nodes),
-        ).tocsr()
-        self.num_nodes = num_nodes
-        self.adj = self.compute_adjlist(self.sparse_adj)
-        self.adj = torch.tensor(self.adj)
-
-    def compute_adjlist(self, sp_adj, max_degree=32):
-        num_data = sp_adj.shape[0]
-        adj = num_data + np.zeros((num_data+1, max_degree), dtype=np.int32)
-
-        for v in range(num_data):
-            neighbors = np.nonzero(sp_adj[v, :])[1]
-            len_neighbors = len(neighbors)
-            if len_neighbors > max_degree:
-                neighbors = np.random.choice(neighbors, max_degree, replace=False)
-                adj[v] = neighbors
-            else:
-                adj[v, :len_neighbors] = neighbors
-
-        return adj
-
-    def from_adjlist(self, adj):
-        u_sampled, index = torch.unique(torch.flatten(adj), return_inverse=True)
-
-        row = (torch.range(0, index.shape[0]-1) / adj.shape[1]).long().to(adj.device)
-        col = index
-        values = torch.ones(index.shape[0]).float().to(adj.device)
-        indices = torch.cat([row.unsqueeze(1), col.unsqueeze(1)], axis=1).t()
-        dense_shape = (adj.shape[0], u_sampled.shape[0])
-
-        support = torch.sparse_coo_tensor(indices, values, dense_shape)
-
-        return support, u_sampled.long()
-
-    def _sample_one_layer(self, x, adj, v, sample_size):
-        support, u = self.from_adjlist(adj)
-
-
-        h_v = torch.sum(torch.matmul(x[v], self.w_s1))
-        h_u = torch.matmul(x[u], self.w_s0)
-        attention = (F.relu(h_v + h_u) + 1) * (1.0 / sample_size)
-        g_u = F.relu(h_u) + 1
-
-        p1 = attention * g_u
-        p1 = p1.cpu()
-
-        if self.num_nodes in u:
-            p1[u == self.num_nodes] = 0
-        p1 = p1 / torch.sum(p1)
-
-        samples = torch.multinomial(p1, sample_size, False)
-        u_sampled = u[samples]
-
-        support_sampled = torch.index_select(support, 1, samples)
-
-        return u_sampled, support_sampled
-
-    def sample(self, x, v, num_layers):
-        all_support = [[] for _ in range(num_layers)]
-        sampled = v
-        x = torch.cat((x, torch.zeros(1, x.shape[1]).to(x.device)), dim=0)
-        for i in range(num_layers - 1, -1, -1):
-            cur_sampled, cur_support = self.sample_one_layer(x, self.adj[sampled], sampled, self.sample_size[i])
-            all_support[i] = cur_support.to(x.device)
-            sampled = cur_sampled
-
-        return x[sampled.to(x.device)], all_support, 0"""
diff --git a/cogdl/datasets/__init__.py b/cogdl/datasets/__init__.py
index 40abd87b..629fbe30 100644
--- a/cogdl/datasets/__init__.py
+++ b/cogdl/datasets/__init__.py
@@ -84,7 +84,7 @@ def build_dataset_from_path(data_path, task):
     "ogbn-arxiv": "cogdl.datasets.ogb",
     "ogbn-products": "cogdl.datasets.ogb",
     "ogbn-proteins": "cogdl.datasets.ogb",
-    "ogbn-mag": "cogdl.datasets.ogb",
+    "ogbn-mag": "cogdl.datasets.pyg_ogb",
     "ogbn-papers100M": "cogdl.datasets.ogb",
     "ogbg-molbace": "cogdl.datasets.ogb",
     "ogbg-molhiv": "cogdl.datasets.ogb",
diff --git a/cogdl/datasets/ogb.py b/cogdl/datasets/ogb.py
index 08e0c284..c6fae035 100644
--- a/cogdl/datasets/ogb.py
+++ b/cogdl/datasets/ogb.py
@@ -31,6 +31,7 @@ def coalesce(row, col, edge_attr=None):
 
 class OGBNDataset(Dataset):
     def __init__(self, root, name):
+        super(OGBNDataset, self).__init__(root)
         dataset = NodePropPredDataset(name, root)
         graph, y = dataset[0]
         x = torch.tensor(graph["node_feat"])
@@ -69,6 +70,12 @@ def get_loss_fn(self):
     def get_evaluator(self):
         return accuracy
 
+    def _download(self):
+        pass
+
+    def _process(self):
+        pass
+
 
 @register_dataset("ogbn-arxiv")
 class OGBArxivDataset(OGBNDataset):
@@ -112,6 +119,7 @@ def __init__(self):
 
 class OGBGDataset(Dataset):
     def __init__(self, root, name):
+        super(OGBGDataset, self).__init__(root)
         self.name = name
         self.dataset = GraphPropPredDataset(self.name, root)
 
@@ -150,6 +158,12 @@ def get_subset(self, subset):
     def get(self, idx):
         return self.graphs[idx]
 
+    def _download(self):
+        pass
+
+    def _process(self):
+        pass
+
     @property
     def num_classes(self):
         return int(self.dataset.num_classes)
diff --git a/cogdl/experiments.py b/cogdl/experiments.py
index d07f5562..cf5b1195 100644
--- a/cogdl/experiments.py
+++ b/cogdl/experiments.py
@@ -64,7 +64,6 @@ 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
 
 
diff --git a/cogdl/match.yml b/cogdl/match.yml
index 07e6ae68..43d99c1b 100644
--- a/cogdl/match.yml
+++ b/cogdl/match.yml
@@ -65,6 +65,7 @@ node_classification:
     - dropedge_gcn
     - unet
     - pprgo
+    - graphsaint
     dataset:
     - cora
     - citeseer
@@ -78,6 +79,7 @@ node_classification:
     - ppi
     - ppi-large
     - reddit
+    - test_small
 unsupervised_node_classification:
   - model:
     - prone
@@ -137,8 +139,6 @@ graph_classification:
     - reddit-b
     - reddit-multi-5k
     - reddit-multi-12k
-    - ogbg-molbace
-    - ogbg-molhiv
 unsupervised_graph_classification:
   - model:
     - infograph
diff --git a/cogdl/models/__init__.py b/cogdl/models/__init__.py
index dc70da0b..8022828e 100644
--- a/cogdl/models/__init__.py
+++ b/cogdl/models/__init__.py
@@ -119,4 +119,5 @@ def build_model(args):
     "pairnorm": "cogdl.models.nn.pairnorm",
     "unsup_graphsage": "cogdl.models.nn.unsup_graphsage",
     "sagpool": "cogdl.models.nn.pyg_sagpool",
+    "graphsaint": "cogdl.models.nn.graphsaint",
 }
diff --git a/cogdl/models/nn/gcn.py b/cogdl/models/nn/gcn.py
index 3789f268..dd2892dc 100644
--- a/cogdl/models/nn/gcn.py
+++ b/cogdl/models/nn/gcn.py
@@ -51,8 +51,8 @@ class TKipfGCN(BaseModel):
     <https://arxiv.org/abs/1609.02907>`_ paper
 
     Args:
-        num_features (int) : Number of input features.
-        num_classes (int) : Number of classes.
+        in_features (int) : Number of input features.
+        out_features (int) : Number of classes.
         hidden_size (int) : The dimension of node representation.
         dropout (float) : Dropout rate for model training.
     """
@@ -78,7 +78,6 @@ def __init__(self, in_feats, hidden_size, out_feats, num_layers, dropout):
         self.layers = nn.ModuleList([GraphConvolution(shapes[i], shapes[i + 1]) for i in range(num_layers)])
         self.num_layers = num_layers
         self.dropout = dropout
-        self.cache = dict()
 
     def get_embeddings(self, x, edge_index):
         edge_index, edge_attr = add_remaining_self_loops(edge_index, num_nodes=x.shape[0])
@@ -90,9 +89,9 @@ def get_embeddings(self, x, edge_index):
             h = self.layers[i](h, edge_index, edge_attr)
         return h
 
-    def forward(self, x, edge_index):
-        edge_index, edge_attr = add_remaining_self_loops(edge_index, num_nodes=x.shape[0])
-        edge_attr = symmetric_normalization(x.shape[0], edge_index, edge_attr)
+    def forward(self, x, edge_index, edge_weight=None):
+        edge_index, edge_weight = add_remaining_self_loops(edge_index, num_nodes=x.shape[0])
+        edge_attr = symmetric_normalization(x.shape[0], edge_index, edge_weight)
         h = x
         for i in range(self.num_layers):
             h = self.layers[i](h, edge_index, edge_attr)
@@ -102,4 +101,7 @@ def forward(self, x, edge_index):
         return h
 
     def predict(self, data):
-        return self.forward(data.x, data.edge_index)
+        if hasattr(data, "norm_aggr"):
+            return self.forward(data.x, data.edge_index, data.norm_aggr)
+        else:
+            return self.forward(data.x, data.edge_index)
diff --git a/cogdl/models/nn/graphsaint.py b/cogdl/models/nn/graphsaint.py
new file mode 100644
index 00000000..600d4991
--- /dev/null
+++ b/cogdl/models/nn/graphsaint.py
@@ -0,0 +1,291 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+import numpy as np
+
+from .. import BaseModel, register_model
+from cogdl.trainers.sampled_trainer import SAINTTrainer
+from cogdl.utils import spmm
+
+
+F_ACT = {"relu": nn.ReLU(), "I": lambda x: x}
+
+"""
+Borrowed from https://github.com/GraphSAINT/GraphSAINT
+"""
+
+
+class HighOrderAggregator(nn.Module):
+    def __init__(self, dim_in, dim_out, dropout=0.0, act="relu", order=1, aggr="mean", bias="norm-nn", **kwargs):
+        """
+        Layer implemented here combines the GraphSAGE-mean [1] layer with MixHop [2] layer.
+        We define the concept of `order`: an order-k layer aggregates neighbor information
+        from 0-hop all the way to k-hop. The operation is approximately:
+            X W_0 [+] A X W_1 [+] ... [+] A^k X W_k
+        where [+] is some aggregation operation such as addition or concatenation.
+
+        Special cases:
+            Order = 0  -->  standard MLP layer
+            Order = 1  -->  standard GraphSAGE layer
+
+        [1]: https://cs.stanford.edu/people/jure/pubs/graphsage-nips17.pdf
+        [2]: https://arxiv.org/abs/1905.00067
+
+        Inputs:
+            dim_in      int, feature dimension for input nodes
+            dim_out     int, feature dimension for output nodes
+            dropout     float, dropout on weight matrices W_0 to W_k
+            act         str, activation function. See F_ACT at the top of this file
+            order       int, see definition above
+            aggr        str, if 'mean' then [+] operation adds features of various hops
+                            if 'concat' then [+] concatenates features of various hops
+            bias        str, if 'bias' then apply a bias vector to features of each hop
+                            if 'norm' then perform batch-normalization on output features
+
+        Outputs:
+            None
+        """
+        super(HighOrderAggregator, self).__init__()
+        assert bias in ["bias", "norm", "norm-nn"]
+        self.order, self.aggr = order, aggr
+        self.act, self.bias = F_ACT[act], bias
+        self.dropout = dropout
+        self.f_lin, self.f_bias = [], []
+        self.offset, self.scale = [], []
+        self.num_param = 0
+        for o in range(self.order + 1):
+            self.f_lin.append(nn.Linear(dim_in, dim_out, bias=False))
+            nn.init.xavier_uniform_(self.f_lin[-1].weight)
+            self.f_bias.append(nn.Parameter(torch.zeros(dim_out)))
+            self.num_param += dim_in * dim_out
+            self.num_param += dim_out
+            self.offset.append(nn.Parameter(torch.zeros(dim_out)))
+            self.scale.append(nn.Parameter(torch.ones(dim_out)))
+            if self.bias == "norm" or self.bias == "norm-nn":
+                self.num_param += 2 * dim_out
+        self.f_lin = nn.ModuleList(self.f_lin)
+        self.f_dropout = nn.Dropout(p=self.dropout)
+        self.params = nn.ParameterList(self.f_bias + self.offset + self.scale)
+        self.f_bias = self.params[: self.order + 1]
+        if self.bias == "norm":
+            self.offset = self.params[self.order + 1 : 2 * self.order + 2]
+            self.scale = self.params[2 * self.order + 2 :]
+        elif self.bias == "norm-nn":
+            final_dim_out = dim_out * ((aggr == "concat") * (order + 1) + (aggr == "mean"))
+            self.f_norm = nn.BatchNorm1d(final_dim_out, eps=1e-9, track_running_stats=True)
+        self.num_param = int(self.num_param)
+
+    def _spmm(self, adj_norm, _feat):
+        """ sparce feature matrix multiply dense feature matrix """
+        # alternative ways: use geometric.propagate or torch.mm
+        return torch.sparse.mm(adj_norm, _feat)
+
+    def _f_feat_trans(self, _feat, _id):
+        feat = self.act(self.f_lin[_id](_feat) + self.f_bias[_id])
+        if self.bias == "norm":
+            mean = feat.mean(dim=1).view(feat.shape[0], 1)
+            var = feat.var(dim=1, unbiased=False).view(feat.shape[0], 1) + 1e-9
+            feat_out = (feat - mean) * self.scale[_id] * torch.rsqrt(var) + self.offset[_id]
+        else:
+            feat_out = feat
+        return feat_out
+
+    def forward(self, inputs):
+        """
+        Inputs:.
+            adj_norm        normalized adj matrix of the subgraph
+            feat_in         2D matrix of input node features
+
+        Outputs:
+            adj_norm        same as input (to facilitate nn.Sequential)
+            feat_out        2D matrix of output node features
+        """
+        x, edge_index, edge_weight = inputs
+        feat_in = self.f_dropout(x)
+        feat_hop = [feat_in]
+        # generate A^i X
+        for o in range(self.order):
+            # propagate(edge_index, x=x, norm=norm)
+            # feat_hop.append(self._spmm(adj_norm, feat_hop[-1]))
+            feat_hop.append(spmm(edge_index, edge_weight, x))
+        feat_partial = [self._f_feat_trans(ft, idf) for idf, ft in enumerate(feat_hop)]
+        if self.aggr == "mean":
+            feat_out = feat_partial[0]
+            for o in range(len(feat_partial) - 1):
+                feat_out += feat_partial[o + 1]
+        elif self.aggr == "concat":
+            feat_out = torch.cat(feat_partial, 1)
+        else:
+            raise NotImplementedError
+        if self.bias == "norm-nn":
+            feat_out = self.f_norm(feat_out)
+        return feat_out, edge_index, edge_weight  # return adj_norm to support Sequential
+
+
+def parse_arch(architecture, aggr, act, bias, hidden_size, num_features):
+    num_layers = len(architecture.split("-"))
+    # set default values, then update by arch_gcn
+    bias_layer = [bias] * num_layers
+    act_layer = [act] * num_layers
+    aggr_layer = [aggr] * num_layers
+    dims_layer = [hidden_size] * num_layers
+    order_layer = [int(order) for order in architecture.split("-")]
+    return [num_features] + dims_layer, order_layer, act_layer, bias_layer, aggr_layer
+
+
+@register_model("graphsaint")
+class GraphSAINT(BaseModel):
+    @staticmethod
+    def add_args(parser):
+        """Add model-specific arguments to the parser."""
+        # fmt: off
+        parser.add_argument("--num-features", type=int)
+        parser.add_argument("--hidden-size", type=int, default=128)
+        parser.add_argument("--architecture", type=str, default="1-1-0")
+        parser.add_argument("--aggr", type=str, default="concat")
+        parser.add_argument("--act", type=str, default="relu")
+        parser.add_argument("--bias", type=str, default="norm")
+        parser.add_argument("--dropout", type=float, default=0.1)
+        parser.add_argument("--weight-decay", type=int, default=0)
+        # fmt: on
+
+    @classmethod
+    def build_model_from_args(cls, args):
+        return cls(
+            args.num_features,
+            args.num_classes,
+            args.architecture,
+            args.aggr,
+            args.act,
+            args.bias,
+            args.weight_decay,
+            args.dropout,
+            args.hidden_size,
+        )
+
+    def __init__(self, num_features, num_classes, architecture, aggr, act, bias, weight_decay, dropout, hidden_size):
+        """
+        Build the multi-layer GNN architecture.
+
+        Inputs:
+            num_classes         int, number of classes a node can belong to
+            arch_gcn            dict, config for each GNN layer
+            train_params        dict, training hyperparameters (e.g., learning rate)
+            feat_full           np array of shape N x f, where N is the total num of
+                                nodes and f is the dimension for input node feature
+            label_full          np array, for single-class classification, the shape
+                                is N x 1 and for multi-class classification, the
+                                shape is N x c (where c = num_classes)
+            cpu_eval            bool, if True, will put the model on CPU.
+
+        Outputs:
+            None
+        """
+        super(GraphSAINT, self).__init__()
+        self.aggregator_cls = HighOrderAggregator
+        self.mulhead = 1
+        self.weight_decay = weight_decay
+        self.dropout = dropout
+        self.sigmoid_loss = True
+        self.num_classes = num_classes
+        self.num_layers = len(architecture.split("-"))
+        _dims, self.order_layer, self.act_layer, self.bias_layer, self.aggr_layer = parse_arch(
+            architecture, aggr, act, bias, hidden_size, num_features
+        )
+        # get layer index for each conv layer, useful for jk net last layer aggregation
+        self.set_idx_conv()
+        self.set_dims(_dims)
+
+        self.loss = 0
+        self.opt_op = None
+
+        # build the model below
+        self.num_params = 0
+        self.aggregators, num_param = self.get_aggregators()
+        self.num_params += num_param
+        self.conv_layers = nn.Sequential(*self.aggregators)
+        self.classifier = HighOrderAggregator(
+            self.dims_feat[-1], self.num_classes, act="I", order=0, dropout=self.dropout, bias="bias"
+        )
+        self.num_params += self.classifier.num_param
+
+    def set_dims(self, dims):
+        """
+        Set the feature dimension / weight dimension for each GNN or MLP layer.
+        We will use the dimensions set here to initialize PyTorch layers.
+
+        Inputs:
+            dims        list, length of node feature for each hidden layer
+
+        Outputs:
+            None
+        """
+        self.dims_feat = [dims[0]] + [
+            ((self.aggr_layer[layer] == "concat") * self.order_layer[layer] + 1) * dims[layer + 1]
+            for layer in range(len(dims) - 1)
+        ]
+        self.dims_weight = [(self.dims_feat[layer], dims[layer + 1]) for layer in range(len(dims) - 1)]
+
+    def set_idx_conv(self):
+        """
+        Set the index of GNN layers for the full neural net. For example, if
+        the full NN is having 1-0-1-0 arch (1-hop graph conv, followed by 0-hop
+        MLP, ...). Then the layer indices will be 0, 2.
+        """
+        idx_conv = np.where(np.array(self.order_layer) >= 1)[0]
+        idx_conv = list(idx_conv[1:] - 1)
+        idx_conv.append(len(self.order_layer) - 1)
+        _o_arr = np.array(self.order_layer)[idx_conv]
+        if np.prod(np.ediff1d(_o_arr)) == 0:
+            self.idx_conv = idx_conv
+        else:
+            self.idx_conv = list(np.where(np.array(self.order_layer) == 1)[0])
+
+    def forward(self, data):
+        if not hasattr(data, "edge_weight"):
+            edge_weight = torch.ones(data.edge_index.shape[1]).to(data.x.device)
+        else:
+            edge_weight = data.edge_weight
+        emb_subg, _, _ = self.conv_layers((data.x, data.edge_index, edge_weight))
+        emb_subg_norm = F.normalize(emb_subg, p=2, dim=1)
+        pred_subg = self.classifier((emb_subg_norm, None, None))[0]
+        return pred_subg
+
+    def _loss(self, preds, labels, norm_loss):
+        """
+        The predictor performs sigmoid (for multi-class) or softmax (for single-class)
+        """
+        if self.sigmoid_loss:
+            norm_loss = norm_loss.unsqueeze(1)
+            return torch.nn.BCEWithLogitsLoss(weight=norm_loss, reduction="sum")(preds, labels)
+        else:
+            _ls = torch.nn.CrossEntropyLoss(reduction="none")(preds, labels)
+            return (norm_loss * _ls).sum()
+
+    def get_aggregators(self):
+        """
+        Return a list of aggregator instances. to be used in self.build()
+        """
+        num_param = 0
+        aggregators = []
+        for layer in range(self.num_layers):
+            aggr = self.aggregator_cls(
+                *self.dims_weight[layer],
+                dropout=self.dropout,
+                act=self.act_layer[layer],
+                order=self.order_layer[layer],
+                aggr=self.aggr_layer[layer],
+                bias=self.bias_layer[layer],
+                mulhead=self.mulhead,
+            )
+            num_param += aggr.num_param
+            aggregators.append(aggr)
+        return aggregators, num_param
+
+    def predict(self, data):
+        return self(data)
+        # return nn.Sigmoid()(preds) if self.sigmoid_loss else F.softmax(preds, dim=1
+
+    @staticmethod
+    def get_trainer(task, args):
+        return SAINTTrainer
diff --git a/cogdl/models/nn/mvgrl.py b/cogdl/models/nn/mvgrl.py
index 7e5d6f27..c3b2c861 100644
--- a/cogdl/models/nn/mvgrl.py
+++ b/cogdl/models/nn/mvgrl.py
@@ -156,9 +156,9 @@ def forward(self, x, edge_index, edge_attr=None):
         idx = np.random.randint(0, adj.shape[-1] - self.sample_size + 1, self.batch_size)
         ba, bd, bf = [], [], []
         for i in idx:
-            ba.append(adj[i : i + self.sample_size, i : i + self.sample_size])
-            bd.append(diff[i : i + self.sample_size, i : i + self.sample_size])
-            bf.append(x[i : i + self.sample_size])
+            ba.append(adj[i: i + self.sample_size, i: i + self.sample_size])
+            bd.append(diff[i: i + self.sample_size, i: i + self.sample_size])
+            bf.append(x[i: i + self.sample_size])
 
         ba = np.array(ba).reshape(self.batch_size, self.sample_size, self.sample_size)
         bd = np.array(bd)
diff --git a/cogdl/models/nn/ppnp.py b/cogdl/models/nn/ppnp.py
index 04483298..85ff6438 100644
--- a/cogdl/models/nn/ppnp.py
+++ b/cogdl/models/nn/ppnp.py
@@ -3,7 +3,6 @@
 
 from .. import BaseModel, register_model
 from cogdl.utils import add_remaining_self_loops, spmm
-
 from .mlp import MLP
 
 
diff --git a/cogdl/models/nn/sgc.py b/cogdl/models/nn/sgc.py
index 491258ed..1a1e87f7 100644
--- a/cogdl/models/nn/sgc.py
+++ b/cogdl/models/nn/sgc.py
@@ -11,12 +11,10 @@ def __init__(self, in_features, out_features, order=3):
         self.in_features = in_features
         self.out_features = out_features
         self.order = order
-        # self.weight = Parameter(torch.FloatTensor(in_features, out_features))
-
         self.W = nn.Linear(in_features, out_features)
 
-    def forward(self, input, edge_index, edge_attr=None):
-        output = self.W(input)
+    def forward(self, x, edge_index, edge_attr=None):
+        output = self.W(x)
         for _ in range(self.order):
             output = spmm(edge_index, edge_attr, output)
         return output
diff --git a/cogdl/tasks/base_task.py b/cogdl/tasks/base_task.py
index f00e3680..c315fe5e 100644
--- a/cogdl/tasks/base_task.py
+++ b/cogdl/tasks/base_task.py
@@ -58,7 +58,7 @@ def set_evaluator(self, dataset):
         self.evaluator = dataset.get_evaluator()
 
     def get_trainer(self, model, args):
-        if hasattr(args, "trainer") and args.trainer:
+        if hasattr(args, "trainer") and args.trainer is not None:
             if "self_auxiliary_task" in args.trainer and not hasattr(model, "get_embeddings"):
                 raise ValueError("Model ({}) must implement get_embeddings method".format(args.model))
             return build_trainer(args)
diff --git a/cogdl/tasks/graph_classification.py b/cogdl/tasks/graph_classification.py
index e73f0a73..68826a1f 100644
--- a/cogdl/tasks/graph_classification.py
+++ b/cogdl/tasks/graph_classification.py
@@ -79,6 +79,7 @@ def __init__(self, args, dataset=None, model=None):
         self.device = "cpu" if not torch.cuda.is_available() or args.cpu else args.device_id[0]
 
         if args.dataset.startswith("ogbg"):
+            self.data = dataset.data
             self.train_loader, self.val_loader, self.test_loader = dataset.get_loader(args)
             model = build_model(args) if model is None else model
         else:
diff --git a/cogdl/tasks/node_classification.py b/cogdl/tasks/node_classification.py
index 8446a762..02eb9cf2 100644
--- a/cogdl/tasks/node_classification.py
+++ b/cogdl/tasks/node_classification.py
@@ -22,6 +22,7 @@ class NodeClassification(BaseTask):
     def add_args(parser: argparse.ArgumentParser):
         """Add task-specific arguments to the parser."""
         # fmt: off
+        parser.add_argument("--missing-rate", type=int, default=0, help="missing rate, from 0 to 100")
         # fmt: on
 
     def __init__(
@@ -67,6 +68,7 @@ def train(self):
             result = self.trainer.fit(self.model, self.dataset)
             if issubclass(type(result), torch.nn.Module):
                 self.model = result
+                self.model.to(self.data.x.device)
             else:
                 return result
         else:
@@ -81,13 +83,14 @@ def train(self):
                 self._train_step()
                 train_acc, _ = self._test_step(split="train")
                 val_acc, val_loss = self._test_step(split="val")
+                print(_, val_loss)
                 epoch_iter.set_description(f"Epoch: {epoch:03d}, Train: {train_acc:.4f}, Val: {val_acc:.4f}")
                 if val_loss <= min_loss or val_acc >= max_score:
                     if val_loss <= best_loss:  # and val_acc >= best_score:
                         best_loss = val_loss
                         best_score = val_acc
                         best_model = copy.deepcopy(self.model)
-                    min_loss = np.min((min_loss, val_loss))
+                    min_loss = np.min((min_loss, val_loss.cpu()))
                     max_score = np.max((max_score, val_acc))
                     patience = 0
                 else:
diff --git a/cogdl/trainers/__init__.py b/cogdl/trainers/__init__.py
index dbefaef1..528c161c 100644
--- a/cogdl/trainers/__init__.py
+++ b/cogdl/trainers/__init__.py
@@ -53,5 +53,6 @@ def build_trainer(args):
     "graphsaint": "cogdl.trainers.sampled_trainer",
     "neighborsampler": "cogdl.trainers.sampled_trainer",
     "clustergcn": "cogdl.trainers.sampled_trainer",
+    "random_partition": "cogdl.trainers.sampled_trainer",
     "self_auxiliary_task": "cogdl.trainers.self_auxiliary_task_trainer",
 }
diff --git a/cogdl/trainers/sampled_trainer.py b/cogdl/trainers/sampled_trainer.py
index d9d02297..6981432e 100644
--- a/cogdl/trainers/sampled_trainer.py
+++ b/cogdl/trainers/sampled_trainer.py
@@ -82,6 +82,7 @@ def add_args(parser: argparse.ArgumentParser):
         parser.add_argument('--num-walks', default=50, type=int, help='number of random walks')
         parser.add_argument('--walk-length', default=20, type=int, help='random walk length')
         parser.add_argument('--size-frontier', default=20, type=int, help='frontier size in multidimensional random walks')
+        parser.add_argument('--valid-cpu', action='store_true', help='run validation on cpu')
         # fmt: on
 
     @classmethod
@@ -90,6 +91,7 @@ def build_trainer_from_args(cls, args):
 
     def __init__(self, args):
         super(SAINTTrainer, self).__init__(args)
+        self.valid_cpu = args.valid_cpu
         self.args_sampler = self.sampler_from_args(args)
 
     def sampler_from_args(self, args):
@@ -103,9 +105,9 @@ def sampler_from_args(self, args):
         }
         return args_sampler
 
-    def fit(self, model: SupervisedModel, dataset: Dataset):
+    def set_data_model(self, dataset: Dataset, model: SupervisedModel):
+        self.dataset = dataset
         self.data = dataset.data
-        self.data.apply(lambda x: x.to(self.device))
         self.model = model.to(self.device)
         self.evaluator = dataset.get_evaluator()
         self.loss_fn = dataset.get_loss_fn()
@@ -118,28 +120,59 @@ def fit(self, model: SupervisedModel, dataset: Dataset):
             self.sampler = RWSampler(self.data, self.args_sampler)
         elif self.args_sampler["sampler"] == "mrw":
             self.sampler = MRWSampler(self.data, self.args_sampler)
+        else:
+            raise NotImplementedError
 
         self.optimizer = torch.optim.Adam(model.parameters(), lr=self.lr, weight_decay=self.weight_decay)
-        best_model = self.train()
-        self.model = best_model
-        return self.model
+
+    def fit(self, model: SupervisedModel, dataset: Dataset):
+        self.set_data_model(dataset, model)
+        return self.train()
 
     def _train_step(self):
-        self.data = self.sampler.get_subgraph("train")
+        self.data = self.sampler.one_batch("train")
         self.data.apply(lambda x: x.to(self.device))
+
+        self.model = self.model.to(self.device)
         self.model.train()
         self.optimizer.zero_grad()
-        self.model.node_classification_loss(self.data).backward()
+
+        mask = self.data.train_mask
+        if len(self.data.y.shape) > 1:
+            logits = self.model.predict(self.data)
+            weight = self.data.norm_loss[mask].unsqueeze(1)
+            loss = torch.nn.BCEWithLogitsLoss(reduction="sum", weight=weight)(logits[mask], self.data.y[mask].float())
+        else:
+            logits = torch.nn.functional.log_softmax(self.model.predict(self.data))
+            loss = (
+                torch.nn.NLLLoss(reduction="none")(logits[mask], self.data.y[mask]) * self.data.norm_loss[mask]
+            ).sum()
+        loss.backward()
         self.optimizer.step()
 
     def _test_step(self, split="val"):
-        self.data = self.sampler.get_subgraph(split)
-        self.data.apply(lambda x: x.to(self.device))
+        self.data = self.sampler.one_batch(split)
+        if split != "train" and self.valid_cpu:
+            self.model = self.model.cpu()
+        else:
+            self.data.apply(lambda x: x.to(self.device))
         self.model.eval()
         masks = {"train": self.data.train_mask, "val": self.data.val_mask, "test": self.data.test_mask}
         with torch.no_grad():
             logits = self.model.predict(self.data)
-        # self.loss_fn(logits[val], self.data.y[val]) * self.data.norm_loss[val]
+
+        # if isinstance(self.dataset, SAINTDataset):
+        #     weight = self.data.norm_loss.unsqueeze(1)
+        #     loss = torch.nn.BCEWithLogitsLoss(reduction="sum", weight=weight)(logits, self.data.y.float())
+        #     metric = multilabel_f1(logits[mask], self.data.y[mask])
+        # else:
+        #     loss = (
+        #         torch.nn.NLLLoss(reduction="none")(F.log_softmax(logits[mask]), self.data.y[mask])
+        #         * self.data.norm_loss[mask]
+        #     ).sum()
+        #     pred = logits[mask].max(1)[1]
+        #     metric = pred.eq(self.data.y[mask]).sum().item() / mask.sum().item()
+
         loss = {key: self.loss_fn(logits[val], self.data.y[val]) for key, val in masks.items()}
         metric = {key: self.evaluator(logits[val], self.data.y[val]) for key, val in masks.items()}
         return metric, loss
@@ -350,89 +383,3 @@ def _test_step(self, split="val"):
         loss = {key: self.loss_fn(logits[val], self.data.y[val]) for key, val in masks.items()}
         metric = {key: self.evaluator(logits[val], self.data.y[val]) for key, val in masks.items()}
         return metric, loss
-
-
-"""
-class LayerSampledTrainer(SampledTrainer):
-    def __init__(self, args):
-        self.device = torch.device('cpu' if args.cpu else 'cuda')
-        self.patience = args.patience
-        self.max_epoch = args.max_epoch
-        self.batch_size = args.batch_size
-
-    def fit(self, model: SamplingNodeClassificationModel, dataset: Dataset):
-        self.model = model.to(self.device)
-        self.data = dataset.data
-        self.data.apply(lambda x: x.to(self.device))
-        self.sampler = LayerSampler(self.data, self.model, {})
-        self.num_nodes = self.data.x.shape[0]
-        self.adj_list = self.data.edge_index.detach().cpu().numpy()
-        self.model.set_adj(self.adj_list, self.num_nodes)
-        self.optimizer = torch.optim.Adam(
-            self.model.parameters(), lr=args.lr, weight_decay=args.weight_decay
-        )
-
-        epoch_iter = tqdm(range(self.max_epoch))
-        patience = 0
-        best_score = 0
-        best_loss = np.inf
-        max_score = 0
-        min_loss = np.inf
-        for epoch in epoch_iter:
-            self._train_step()
-            train_acc, _ = self._test_step(split="train")
-            val_acc, val_loss = self._test_step(split="val")
-            epoch_iter.set_description(
-                f"Epoch: {epoch:03d}, Train: {train_acc:.4f}, Val: {val_acc:.4f}"
-            )
-            if val_loss <= min_loss or val_acc >= max_score:
-                if val_loss <= best_loss:  # and val_acc >= best_score:
-                    best_loss = val_loss
-                    best_score = val_acc
-                    best_model = copy.deepcopy(self.model)
-                min_loss = np.min((min_loss, val_loss))
-                max_score = np.max((max_score, val_acc))
-                patience = 0
-            else:
-                patience += 1
-                if patience == self.patience:
-                    self.model = best_model
-                    epoch_iter.close()
-                    break
-
-    def _train_step(self):
-        self.model.train()
-        train_nodes = np.where(self.data.train_mask.detach().cpu().numpy())[0]
-        train_labels = self.data.y.detach().cpu().numpy()
-        for batch_nodes, batch_labels in get_batches(train_nodes, train_labels, batch_size=self.batch_size):
-            batch_nodes = torch.LongTensor(batch_nodes)
-            batch_labels = torch.LongTensor(batch_labels).to(self.device)
-            sampled_x, sampled_adj, var_loss = self.sampler.sampling(self.data.x, batch_nodes)
-            self.optimizer.zero_grad()
-            output = self.model(sampled_x, sampled_adj)
-            loss = F.nll_loss(output, batch_labels) + 0.5 * var_loss
-            loss.backward()
-            self.optimizer.step()
-
-    def _test_step(self, split="val"):
-        self.model.eval()
-        _, mask = list(self.data(f"{split}_mask"))[0]
-        test_nodes = np.where(mask.detach().cpu().numpy())[0]
-        test_labels = self.data.y.detach().cpu().numpy()
-        all_loss = []
-        all_acc = []
-        for batch_nodes, batch_labels in get_batches(test_nodes, test_labels, batch_size=self.batch_size):
-            batch_nodes = torch.LongTensor(batch_nodes)
-            batch_labels = torch.LongTensor(batch_labels).to(self.device)
-            sampled_x, sampled_adj, var_loss = self.model.sampling(self.data.x, batch_nodes)
-            with torch.no_grad():
-                logits = self.model(sampled_x, sampled_adj)
-                loss = F.nll_loss(logits, batch_labels)
-            pred = logits.max(1)[1]
-            acc = pred.eq(self.data.y[batch_nodes]).sum().item() / batch_nodes.shape[0]
-
-            all_loss.append(loss.item())
-            all_acc.append(acc)
-
-        return np.mean(all_acc), np.mean(all_loss)
-"""
diff --git a/cogdl/utils/evaluator.py b/cogdl/utils/evaluator.py
index 2fa58bce..0427ad14 100644
--- a/cogdl/utils/evaluator.py
+++ b/cogdl/utils/evaluator.py
@@ -22,6 +22,7 @@ def multilabel_f1(y_pred, y_true, sigmoid=False):
         y_pred[y_pred > 0] = 1
         y_pred[y_pred <= 0] = 0
     tp = (y_true * y_pred).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)
 


Test patch:
diff --git a/cogdl/datasets/test_data.py b/cogdl/datasets/test_data.py
index da9fc801..c8c7c4fc 100644
--- a/cogdl/datasets/test_data.py
+++ b/cogdl/datasets/test_data.py
@@ -30,7 +30,7 @@ def __init__(self):
         self.data.test_mask = torch.tensor(
             [False, False, True, False, True, False, False, True, True, False, True, False]
         )
-        self.num_classes = 4
+        # self.num_classes = 4
         self.transform = None
 
     def get(self, idx):
diff --git a/tests/tasks/test_node_classification.py b/tests/tasks/test_node_classification.py
index 26fad461..c1572877 100644
--- a/tests/tasks/test_node_classification.py
+++ b/tests/tasks/test_node_classification.py
@@ -254,6 +254,7 @@ def test_gcn_cora_sampler():
     args.dataset = "cora"
     args.trainer = "graphsaint"
     args.model = "gcn"
+    args.valid_cpu = True
     args.cpu = True
     args.num_layers = 2
     args.sample_coverage = 20
@@ -270,6 +271,29 @@ def test_gcn_cora_sampler():
         assert 0 <= ret["Acc"] <= 1
 
 
+def test_graphsaint_cora():
+    args = get_default_args()
+    args.task = "node_classification"
+    args.dataset = "cora"
+    args.trainer = "graphsaint"
+    args.model = "graphsaint"
+    args.valid_cpu = True
+    args.cpu = True
+    args.architecture = "1-1-0"
+    args.aggr = "concat"
+    args.act = "relu"
+    args.bias = "norm"
+    args.sample_coverage = 10
+    args.size_subgraph = 200
+    args.num_walks = 20
+    args.walk_length = 10
+    args.size_frontier = 20
+    args.sampler = "node"
+    task = build_task(args)
+    ret = task.train()
+    assert 0 <= ret["Acc"] <= 1
+
+
 def test_unet_citeseer():
     args = get_default_args()
     args.cpu = True
@@ -701,6 +725,7 @@ def test_gcn_ppi():
     test_deepergcn_cora()
     test_grand_cora()
     test_gcn_cora_sampler()
+    test_graphsaint_cora()
     test_gpt_gnn_cora()
     test_sign_cora()
     test_jknet_jknet_cora()