Repository: THUDM/CogDL (language: python)

Issue:
Refactor Data, Dataset
Hi,

Our current implementations of data preprocessing and data loading are borrowed from pyg. This part needs refactor before release.


Discussion / hints:

Fix patch:
diff --git a/cogdl/tasks/unsupervised_node_classification.py b/cogdl/tasks/unsupervised_node_classification.py
index 0c389adf..72e9c190 100644
--- a/cogdl/tasks/unsupervised_node_classification.py
+++ b/cogdl/tasks/unsupervised_node_classification.py
@@ -40,6 +40,7 @@ def __init__(self, args, dataset=None, model=None):
         super(UnsupervisedNodeClassification, self).__init__(args)
         dataset = build_dataset(args) if dataset is None else dataset
 
+        self.dataset = dataset
         self.data = dataset[0]
 
         self.num_nodes = self.data.y.shape[0]
@@ -103,7 +104,7 @@ def save_emb(self, embs):
 
     def train(self):
         if self.trainer is not None:
-            return self.trainer.fit(self.model, self.data)
+            return self.trainer.fit(self.model, self.dataset)
         if self.load_emb_path is None:
             if "gcc" in self.model_name:
                 features_matrix = self.model.train(self.data)
diff --git a/cogdl/trainers/__init__.py b/cogdl/trainers/__init__.py
index ea9c04e0..e85dd3a4 100644
--- a/cogdl/trainers/__init__.py
+++ b/cogdl/trainers/__init__.py
@@ -54,6 +54,7 @@ def build_trainer(args):
     "neighborsampler": "cogdl.trainers.sampled_trainer",
     "clustergcn": "cogdl.trainers.sampled_trainer",
     "random_cluster": "cogdl.trainers.sampled_trainer",
+    "self_supervised": "cogdl.trainers.self_supervised_trainer",
     "self_auxiliary_task_pretrain": "cogdl.trainers.self_auxiliary_task_trainer",
     "self_auxiliary_task_joint": "cogdl.trainers.self_auxiliary_task_trainer",
     "m3s": "cogdl.trainers.m3s_trainer",
diff --git a/cogdl/trainers/self_auxiliary_task_trainer.py b/cogdl/trainers/self_auxiliary_task_trainer.py
index 23245456..a45dd429 100644
--- a/cogdl/trainers/self_auxiliary_task_trainer.py
+++ b/cogdl/trainers/self_auxiliary_task_trainer.py
@@ -15,18 +15,21 @@
     SupervisedHomogeneousNodeClassificationModel,
 )
 from cogdl.trainers.supervised_model_trainer import SupervisedHomogeneousNodeClassificationTrainer
+from cogdl.utils import dropout_adj
 from . import register_trainer
 from .self_supervised_trainer import LogRegTrainer
 
 from sklearn.metrics.cluster import normalized_mutual_info_score
+from sklearn.cluster import KMeans
 
 
 class SSLTask:
-    def __init__(self, edge_index, features, device):
-        self.edge_index = edge_index
-        self.num_nodes = int(max(edge_index[0].max(), edge_index[1].max()).cpu()) + 1
-        self.num_edges = edge_index[0].shape[0]
-        self.features = features
+    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
         self.device = device
         self.cached_edges = None
 
@@ -38,8 +41,8 @@ def make_loss(self, embeddings):
 
 
 class EdgeMask(SSLTask):
-    def __init__(self, edge_index, features, hidden_size, mask_ratio, device):
-        super().__init__(edge_index, features, device)
+    def __init__(self, graph, hidden_size, mask_ratio, device):
+        super().__init__(graph, device)
         self.linear = nn.Linear(hidden_size, 2).to(device)
         self.mask_ratio = mask_ratio
 
@@ -56,9 +59,10 @@ def transform_data(self):
             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)
+            self.node_pairs = torch.cat([self.masked_edges, self.neg_edges], 1).to(self.device)
+            self.graph.edge_index = self.cached_edges
 
-        return self.cached_edges, self.features
+        return self.graph.to(self.device)
 
     def make_loss(self, embeddings):
         embeddings = self.linear(torch.abs(embeddings[self.node_pairs[0]] - embeddings[self.node_pairs[1]]))
@@ -70,7 +74,7 @@ def neg_sample(self, edge_num):
         exclude = set([(_[0], _[1]) for _ in list(edges)])
         itr = self.sample(exclude)
         sampled = [next(itr) for _ in range(edge_num)]
-        return torch.tensor(sampled).t().to(self.device)
+        return torch.tensor(sampled).t()
 
     def sample(self, exclude):
         while True:
@@ -82,9 +86,9 @@ def sample(self, exclude):
 
 
 class AttributeMask(SSLTask):
-    def __init__(self, edge_index, features, hidden_size, train_mask, mask_ratio, device):
-        super().__init__(edge_index, features, device)
-        self.linear = nn.Linear(hidden_size, features.shape[1]).to(device)
+    def __init__(self, graph, hidden_size, train_mask, mask_ratio, device):
+        super().__init__(graph, 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
@@ -95,10 +99,11 @@ def transform_data(self):
             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]).to(self.device)
-            self.pseudo_labels = self.features[self.masked_nodes]
+            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.edge_index, self.cached_features
+        return self.graph.to(self.device)
 
     def make_loss(self, embeddings):
         embeddings = self.linear(embeddings[self.masked_nodes])
@@ -107,26 +112,27 @@ def make_loss(self, embeddings):
 
 
 class PairwiseDistance(SSLTask):
-    def __init__(self, edge_index, features, hidden_size, class_split, sampling, device):
-        super().__init__(edge_index, features, device)
+    def __init__(self, graph, hidden_size, class_split, sampling, dropedge_rate, num_centers, device):
+        super().__init__(graph, device)
         self.nclass = len(class_split) + 1
         self.class_split = class_split
         self.max_distance = self.class_split[self.nclass - 2][1]
         self.sampling = sampling
-        self.batch_size = 256
+        self.dropedge_rate = dropedge_rate
+        self.num_centers = num_centers
         self.linear = nn.Linear(hidden_size, self.nclass).to(device)
         self.get_distance()
 
     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.batch_size)
+            node_idx = random.sample(range(self.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),
             ).tocsr()
 
-            num_samples = tqdm(range(self.batch_size))
+            num_samples = tqdm(range(self.num_centers))
             for i in num_samples:
                 num_samples.set_description(f"Generating node pairs {i:03d}")
                 idx = node_idx[i]
@@ -173,8 +179,11 @@ 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:
-                edge_index = torch.stack(self.edge_index).cpu().numpy().transpose()
-                self.dis_node_pairs[0] = edge_index
+                self.dis_node_pairs[0] = torch.stack(self.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())
@@ -198,39 +207,44 @@ def get_distance(self):
             self.dis_node_pairs.append(tmp)
 
     def transform_data(self):
-        return self.edge_index, self.features
+        self.graph.edge_index, _ = dropout_adj(edge_index=self.edge_index, drop_rate=self.dropedge_rate)
+        return self.graph.to(self.device)
 
-    def make_loss(self, embeddings, k=4000):
-        node_pairs, pseudo_labels = self.sample(k)
+    def make_loss(self, embeddings, sample=True, k=4000):
+        node_pairs, pseudo_labels = self.sample(sample, k)
         # print(node_pairs, pseudo_labels)
         embeddings = self.linear(torch.abs(embeddings[node_pairs[0]] - embeddings[node_pairs[1]]))
         output = F.log_softmax(embeddings, dim=1)
         return F.nll_loss(output, pseudo_labels)
 
-    def sample(self, k):
+    def sample(self, sample, k):
         sampled = torch.tensor([]).long()
         pseudo_labels = torch.tensor([]).long()
         for i in range(self.nclass):
             tmp = self.dis_node_pairs[i]
-            x = int(random.random() * (len(tmp) - k))
-            sampled = torch.cat([sampled, torch.tensor(tmp[x : x + k]).long().t()], 1)
-            """
-            indices = np.random.choice(np.arange(len(tmp)), k, replace=False)
-            sampled = torch.cat([sampled, torch.tensor(tmp[indices]).long().t()], 1)
-            """
-            pseudo_labels = torch.cat([pseudo_labels, torch.ones(k).long() * i])
+            if sample:
+                x = int(random.random() * (len(tmp) - k))
+                sampled = torch.cat([sampled, torch.tensor(tmp[x : x + k]).long().t()], 1)
+                """
+                indices = np.random.choice(np.arange(len(tmp)), k, replace=False)
+                sampled = torch.cat([sampled, torch.tensor(tmp[indices]).long().t()], 1)
+                """
+                pseudo_labels = torch.cat([pseudo_labels, torch.ones(k).long() * i])
+            else:
+                sampled = torch.cat([sampled, torch.tensor(tmp).long().t()], 1)
+                pseudo_labels = torch.cat([pseudo_labels, torch.ones(len(tmp)).long() * i])
         return sampled.to(self.device), pseudo_labels.to(self.device)
 
 
 class Distance2Clusters(SSLTask):
-    def __init__(self, edge_index, features, hidden_size, num_clusters, device):
-        super().__init__(edge_index, features, device)
+    def __init__(self, graph, hidden_size, num_clusters, device):
+        super().__init__(graph, device)
         self.num_clusters = num_clusters
         self.linear = nn.Linear(hidden_size, num_clusters).to(device)
         self.gen_cluster_info()
 
     def transform_data(self):
-        return self.edge_index, self.features
+        return self.graph.to(self.device)
 
     def gen_cluster_info(self, use_metis=False):
         G = nx.Graph()
@@ -240,14 +254,6 @@ def gen_cluster_info(self, use_metis=False):
 
             _, parts = metis.part_graph(G, self.num_clusters)
         else:
-            """
-            from sklearn.cluster import SpectralClustering
-
-            clustering = SpectralClustering(
-                n_clusters=self.num_clusters, assign_labels="discretize", random_state=0
-            ).fit(self.features.cpu())
-            parts = clustering.labels_
-            """
             from sklearn.cluster import KMeans
 
             clustering = KMeans(n_clusters=self.num_clusters, random_state=0).fit(self.features.cpu())
@@ -276,8 +282,8 @@ def make_loss(self, embeddings):
 
 
 class PairwiseAttrSim(SSLTask):
-    def __init__(self, edge_index, features, hidden_size, k, device):
-        super().__init__(edge_index, features, device)
+    def __init__(self, graph, hidden_size, k, device):
+        super().__init__(graph, device)
         self.k = k
         self.linear = nn.Linear(hidden_size, 1).to(self.device)
         self.get_attr_sim()
@@ -319,7 +325,7 @@ def get_avg_distance(self, idx_sorted, k, sampled):
     def get_attr_sim(self):
         from sklearn.metrics.pairwise import cosine_similarity
 
-        sims = cosine_similarity(self.features.detach().cpu().numpy())
+        sims = cosine_similarity(self.features.numpy())
         idx_sorted = sims.argsort(1)
         self.node_pairs = None
         self.pseudo_labels = None
@@ -345,7 +351,7 @@ def sample(self, k, num_nodes):
         return np.array(sampled)
 
     def transform_data(self):
-        return self.edge_index, self.features
+        return self.graph.to(self.device)
 
     def make_loss(self, embeddings):
         """
@@ -359,9 +365,9 @@ def make_loss(self, embeddings):
 
 
 class Distance2ClustersPP(SSLTask):
-    def __init__(self, edge_index, features, labels, hidden_size, num_clusters, k, device):
-        super().__init__(edge_index, features, device)
-        self.labels = labels.cpu().numpy()
+    def __init__(self, graph, labels, hidden_size, num_clusters, k, device):
+        super().__init__(graph, device)
+        self.labels = labels.numpy()
         self.k = k
         self.num_clusters = num_clusters
         self.clusters = None
@@ -408,7 +414,7 @@ def update_cluster(self):
         self.pseudo_labels = self.pseudo_labels.float().to(self.device)
 
     def transform_data(self):
-        return self.edge_index, self.features
+        return self.graph.to(self.device)
 
     def make_loss(self, embeddings):
         embeddings = self.linear(embeddings)
@@ -431,6 +437,7 @@ def __init__(self, args):
         self.hidden_size = args.hidden_size
         self.label_mask = args.label_mask
         self.sampling = args.sampling
+        self.dropedge_rate = args.dropedge_rate
         self.mask_ratio = args.mask_ratio
 
     def resplit_data(self, data):
@@ -449,28 +456,25 @@ def set_loss_eval(self, dataset):
 
     def set_agent(self):
         if self.auxiliary_task == "edgemask":
-            self.agent = EdgeMask(self.data.edge_index, self.data.x, self.hidden_size, self.mask_ratio, self.device)
+            self.agent = EdgeMask(self.data, self.hidden_size, self.mask_ratio, self.device)
         elif self.auxiliary_task == "attributemask":
-            self.agent = AttributeMask(
-                self.data.edge_index, self.data.x, self.hidden_size, self.data.train_mask, self.mask_ratio, self.device
-            )
+            self.agent = AttributeMask(self.data, self.hidden_size, self.data.train_mask, self.mask_ratio, self.device)
         elif self.auxiliary_task == "pairwise-distance":
             self.agent = PairwiseDistance(
-                self.data.edge_index,
-                self.data.x,
+                self.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(self.data.edge_index, self.data.x, self.hidden_size, 30, self.device)
+            self.agent = Distance2Clusters(self.data, self.hidden_size, 30, self.device)
         elif self.auxiliary_task == "pairwise-attr-sim":
-            self.agent = PairwiseAttrSim(self.data.edge_index, self.data.x, self.hidden_size, 5, self.device)
+            self.agent = PairwiseAttrSim(self.data, self.hidden_size, 5, self.device)
         elif self.auxiliary_task == "distance2clusters++":
-            self.agent = Distance2ClustersPP(
-                self.data.edge_index, self.data.x, self.data.y, self.hidden_size, 5, 1, self.device
-            )
+            self.agent = Distance2ClustersPP(self.data, self.data.y, self.hidden_size, 5, 1, self.device)
         else:
             raise Exception(
                 "auxiliary task must be edgemask, pairwise-distance, distance2clusters, distance2clusters++ or pairwise-attr-sim"
@@ -500,20 +504,25 @@ def add_args(parser: argparse.ArgumentParser):
         parser.add_argument('--auxiliary-task', default="none", type=str)
         parser.add_argument('--label-mask', default=0, type=float)
         parser.add_argument("--mask-ratio", default=0.1, type=float)
+        parser.add_argument("--dropedge-rate", default=0, type=float)
+        parser.add_argument('--alpha', default=10, type=float)
         parser.add_argument('--sampling', action="store_true")
         parser.add_argument("--freeze", action="store_true")
+        parser.add_argument("--agc-eval", action="store_true")
         # fmt: on
 
     def __init__(self, args):
         super().__init__(args)
+        self.alpha = args.alpha
         self.freeze = args.freeze
+        self.agc_eval = args.agc_eval
 
     def fit(self, model: SupervisedHomogeneousNodeClassificationModel, dataset: Dataset):
-        # self.resplit_data(dataset.data)
         self.data = dataset.data
-        self.original_data = dataset.data
-        self.data.to(self.device)
+        self.data.add_remaining_self_loops()
         self.set_agent()
+        self.original_data = dataset.data
+        self.original_data.add_remaining_self_loops()
         self.model = model
         self.set_loss_eval(dataset)
 
@@ -548,11 +557,12 @@ def finetune(self):
 
         embeddings = self.best_model.get_embeddings(self.original_data).detach()
         nclass = int(torch.max(self.data.y) + 1)
-        """
-        kmeans = KMeans(n_clusters=nclass, random_state=0).fit(embeddings.cpu().numpy())
-        clusters = kmeans.labels_
-        print("cluster NMI: %.4lf" % (normalized_mutual_info_score(clusters, self.data.y)))
-        """
+
+        if self.agc_eval:
+            kmeans = KMeans(n_clusters=nclass, random_state=0).fit(embeddings.cpu().numpy())
+            clusters = kmeans.labels_
+            print("cluster NMI: %.4lf" % (normalized_mutual_info_score(clusters, self.data.y.cpu())))
+
         if self.freeze:
             opt = {
                 "idx_train": self.original_data.train_mask,
@@ -560,7 +570,7 @@ def finetune(self):
                 "idx_test": self.original_data.test_mask,
                 "num_classes": nclass,
             }
-            result = LogRegTrainer().train(embeddings, self.original_data.y, opt)
+            result = LogRegTrainer().train(embeddings, self.original_data.y, opt, self.loss_fn, self.evaluator)
             print(f"TestAcc: {result: .4f}")
             return dict(Acc=result)
         else:
@@ -586,14 +596,14 @@ def finetune(self):
 
     def _pretrain_step(self):
         with self.data.local_graph():
-            self.data.edge_index, self.data.x = self.agent.transform_data()
+            self.data = self.agent.transform_data()
             self.model.train()
             self.optimizer.zero_grad()
             embeddings = self.model.get_embeddings(self.data)
-            loss = self.agent.make_loss(embeddings)
+            loss = self.alpha * self.agent.make_loss(embeddings)
         loss.backward()
         self.optimizer.step()
-        return loss.item()
+        return loss.item() / self.alpha
 
     def _train_step(self):
         self.model.train()
@@ -613,6 +623,7 @@ def add_args(parser: argparse.ArgumentParser):
         parser.add_argument('--alpha', default=10, type=float)
         parser.add_argument('--label-mask', default=0, type=float)
         parser.add_argument("--mask-ratio", default=0.1, type=float)
+        parser.add_argument("--dropedge-rate", default=0, type=float)
         parser.add_argument('--sampling', action="store_true")
         # fmt: on
 
@@ -626,8 +637,8 @@ def fit(self, model: SupervisedHomogeneousNodeClassificationModel, dataset: Data
         self.original_data = dataset.data
         self.data.to(self.device)
         self.set_agent()
-        self.data.edge_index, self.data.x = self.agent.transform_data()
-        self.original_data = self.original_data.to(self.device)
+        self.data = self.agent.transform_data()
+        self.original_data.apply(lambda x: x.to(self.device))
         self.model = model
         self.set_loss_eval(dataset)
 
@@ -645,12 +656,12 @@ def fit(self, model: SupervisedHomogeneousNodeClassificationModel, dataset: Data
         for epoch in epoch_iter:
             if self.auxiliary_task == "distance2clusters++" and epoch % 40 == 0:
                 self.agent.update_cluster()
-            self._train_step()
+            aux_loss = self._train_step()
             train_acc, _ = self._test_step(split="train")
             val_acc, val_loss = self._test_step(split="val")
             test_acc, test_loss = self._test_step(split="test")
             epoch_iter.set_description(
-                f"Epoch: {epoch:03d}, Train: {train_acc:.4f}, Val: {val_acc:.4f}, Test: {test_acc:.4f}"
+                f"Epoch: {epoch:03d}, Train: {train_acc:.4f}, Val: {val_acc:.4f}, Test: {test_acc:.4f}, Aux loss: {aux_loss:.4f}"
             )
             if val_loss <= min_loss or val_acc >= max_score:
                 if val_loss <= best_loss:  # and val_acc >= best_score:
@@ -665,10 +676,13 @@ def fit(self, model: SupervisedHomogeneousNodeClassificationModel, dataset: Data
 
     def _train_step(self):
         with self.data.local_graph():
-            self.data.edge_index, self.data.x = self.agent.transform_data()
+            self.data = self.agent.transform_data()
             self.model.train()
             self.optimizer.zero_grad()
             embeddings = self.model.get_embeddings(self.data)
             loss = self.model.node_classification_loss(self.data) + self.alpha * self.agent.make_loss(embeddings)
+            aux_loss = self.agent.make_loss(embeddings)
         loss.backward()
         self.optimizer.step()
+
+        return aux_loss
diff --git a/cogdl/trainers/self_supervised_trainer.py b/cogdl/trainers/self_supervised_trainer.py
index 8e08fde4..7245fc62 100644
--- a/cogdl/trainers/self_supervised_trainer.py
+++ b/cogdl/trainers/self_supervised_trainer.py
@@ -1,12 +1,15 @@
 import os
+import argparse
 from tqdm import tqdm
 import numpy as np
 
 import torch
 from torch import nn
 from .base_trainer import BaseTrainer
+from . import register_trainer
 
 
+@register_trainer("self_supervised")
 class SelfSupervisedTrainer(BaseTrainer):
     def __init__(self, args):
         super(SelfSupervisedTrainer, self).__init__()
@@ -18,12 +21,24 @@ def __init__(self, args):
         self.save_dir = args.save_dir
         self.load_emb_path = args.load_emb_path
         self.lr = args.lr
+        self.sampling = args.sampling
+        self.sample_size = args.sample_size
+
+    @staticmethod
+    def add_args(parser: argparse.ArgumentParser):
+        """Add trainer-specific arguments to the parser."""
+        # fmt: off
+        parser.add_argument("--sampling", action="store_true")
+        parser.add_argument("--sample-size", type=int, default=20000)
+        # fmt: on
 
     @classmethod
     def build_trainer_from_args(cls, args):
         return cls(args)
 
-    def fit(self, model, data):
+    def fit(self, model, dataset):
+        data = dataset.data
+        data.add_remaining_self_loops()
         self.data = data
 
         if self.load_emb_path is not None:
@@ -31,8 +46,6 @@ def fit(self, model, data):
             embeds = torch.from_numpy(embeds).to(self.device)
             return self.evaluate(embeds)
 
-        self.data.to(self.device)
-
         best = 1e9
         cnt_wait = 0
         optimizer = torch.optim.Adam(model.parameters(), lr=self.lr, weight_decay=0.0)
@@ -42,31 +55,39 @@ def fit(self, model, data):
 
         model.train()
         for epoch in epoch_iter:
-            optimizer.zero_grad()
+            with self.data.local_graph():
+                if self.sampling:
+                    idx = np.random.choice(np.arange(self.data.num_nodes), self.sample_size, replace=False)
+                    self.data = data.subgraph(idx)
 
-            loss = model.node_classification_loss(data)
-            epoch_iter.set_description(f"Epoch: {epoch:03d}, Loss: {loss.item(): .4f}")
+                self.data.to(self.device)
+                optimizer.zero_grad()
 
-            if loss < best:
-                best = loss
-                cnt_wait = 0
-            else:
-                cnt_wait += 1
+                loss = model.node_classification_loss(self.data)
+                epoch_iter.set_description(f"Epoch: {epoch:03d}, Loss: {loss.item(): .4f}")
 
-            if cnt_wait == self.patience:
-                print("Early stopping!")
-                break
+                if loss < best:
+                    best = loss
+                    cnt_wait = 0
+                else:
+                    cnt_wait += 1
 
-            loss.backward()
-            optimizer.step()
+                if cnt_wait == self.patience:
+                    print("Early stopping!")
+                    break
 
+                loss.backward()
+                optimizer.step()
+
+        self.data = data
+        self.data.to(self.device)
         with torch.no_grad():
-            embeds = model.embed(data)
+            embeds = model.embed(self.data)
         self.save_embed(embeds)
 
-        return self.evaluate(embeds)
+        return self.evaluate(embeds, dataset.get_loss_fn(), dataset.get_evaluator())
 
-    def evaluate(self, embeds):
+    def evaluate(self, embeds, loss_fn=None, evaluator=None):
         nclass = int(torch.max(self.data.y) + 1)
         opt = {
             "idx_train": self.data.train_mask.to(self.device),
@@ -74,7 +95,7 @@ def evaluate(self, embeds):
             "idx_test": self.data.test_mask.to(self.device),
             "num_classes": nclass,
         }
-        result = LogRegTrainer().train(embeds, self.data.y.to(self.device), opt)
+        result = LogRegTrainer().train(embeds, self.data.y.to(self.device), opt, loss_fn, evaluator)
         print(f"TestAcc: {result: .4f}")
         return dict(Acc=result)
 
@@ -105,7 +126,7 @@ def forward(self, seq):
 
 
 class LogRegTrainer(object):
-    def train(self, data, labels, opt):
+    def train(self, data, labels, opt, loss_fn=None, evaluator=None):
         device = data.device
         idx_train = opt["idx_train"].to(device)
         idx_test = opt["idx_test"].to(device)
@@ -120,7 +141,7 @@ def train(self, data, labels, opt):
         test_lbls = labels[idx_test]
         tot = 0
 
-        xent = nn.CrossEntropyLoss()
+        xent = nn.CrossEntropyLoss() if loss_fn is None else loss_fn
 
         for _ in range(50):
             log = LogReg(nhid, nclass).to(device)
@@ -138,7 +159,10 @@ def train(self, data, labels, opt):
                 optimizer.step()
 
             logits = log(test_embs)
-            preds = torch.argmax(logits, dim=1)
-            acc = torch.sum(preds == test_lbls).float() / test_lbls.shape[0]
-            tot += acc.item()
+            if evaluator is None:
+                preds = torch.argmax(logits, dim=1)
+                acc = torch.sum(preds == test_lbls).float() / test_lbls.shape[0]
+            else:
+                acc = evaluator(logits, test_lbls)
+            tot += acc
         return tot / 50


Test patch:
diff --git a/tests/models/ssl/test_contrastive_models.py b/tests/models/ssl/test_contrastive_models.py
index 2ad3842b..0120a82d 100644
--- a/tests/models/ssl/test_contrastive_models.py
+++ b/tests/models/ssl/test_contrastive_models.py
@@ -19,6 +19,8 @@ def get_default_args():
         "training_percents": [0.1],
         "activation": "relu",
         "residual": False,
+        "sampling": False,
+        "sample_size": 20,
         "norm": None,
     }
     return build_args_from_dict(default_dict)
@@ -43,6 +45,8 @@ def get_unsupervised_nn_args():
         "task": "unsupervised_node_classification",
         "checkpoint": False,
         "load_emb_path": None,
+        "sampling": False,
+        "sample_size": 20,
         "training_percents": [0.1],
     }
     return build_args_from_dict(default_dict)
diff --git a/tests/models/ssl/test_generative_models.py b/tests/models/ssl/test_generative_models.py
index 90b20596..2546c2de 100644
--- a/tests/models/ssl/test_generative_models.py
+++ b/tests/models/ssl/test_generative_models.py
@@ -31,6 +31,8 @@ def get_default_args():
         "checkpoint": False,
         "num_layers": 2,
         "activation": "relu",
+        "dropedge_rate": 0,
+        "agc_eval": False,
         "residual": False,
         "norm": None,
     }
@@ -188,7 +190,7 @@ def test_m3s():
 
 
 if __name__ == "__main__":
-    test_supergat()
+    # test_supergat()
     test_m3s()
     test_edgemask()
     test_edgemask_pt_ft()
diff --git a/tests/tasks/test_encode_paper.py b/tests/tasks/test_encode_paper.py
index dfdb2131..f6a101c9 100644
--- a/tests/tasks/test_encode_paper.py
+++ b/tests/tasks/test_encode_paper.py
@@ -2,24 +2,30 @@
 
 
 def test_encode_paper():
-    tokenizer, model = oagbert("oagbert-v2")
-    title = 'BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding'
-    abstract = 'We introduce a new language representation model called BERT, which stands for Bidirectional Encoder Representations from Transformers. Unlike recent language representation...'
-    authors = ['Jacob Devlin', 'Ming-Wei Chang', 'Kenton Lee', 'Kristina Toutanova']
-    venue = 'north american chapter of the association for computational linguistics'
-    affiliations = ['Google']
-    concepts = ['language model', 'natural language inference', 'question answering']
+    tokenizer, model = oagbert("oagbert-v2-test")
+    title = "BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding"
+    abstract = "We introduce a new language representation model called BERT, which stands for Bidirectional Encoder Representations from Transformers. Unlike recent language representation..."
+    authors = ["Jacob Devlin", "Ming-Wei Chang", "Kenton Lee", "Kristina Toutanova"]
+    venue = "north american chapter of the association for computational linguistics"
+    affiliations = ["Google"]
+    concepts = ["language model", "natural language inference", "question answering"]
     # encode paper
     paper_info = model.encode_paper(
-        title=title, abstract=abstract, venue=venue, authors=authors, concepts=concepts, affiliations=affiliations, reduction="max"
+        title=title,
+        abstract=abstract,
+        venue=venue,
+        authors=authors,
+        concepts=concepts,
+        affiliations=affiliations,
+        reduction="max",
     )
 
     assert len(paper_info) == 5
-    assert paper_info['text'][0]['type'] == 'TEXT'
-    assert len(paper_info['authors']) == 4
-    assert len(paper_info['venue'][0]['token_ids']) == 9
-    assert tuple(paper_info['text'][0]['sequence_output'].shape) == (43, 768)
-    assert len(paper_info['text'][0]['pooled_output']) == 768
+    assert paper_info["text"][0]["type"] == "TEXT"
+    assert len(paper_info["authors"]) == 4
+    assert len(paper_info["venue"][0]["token_ids"]) == 9
+    assert tuple(paper_info["text"][0]["sequence_output"].shape) == (43, 768)
+    assert len(paper_info["text"][0]["pooled_output"]) == 768
 
 
 if __name__ == "__main__":