[Dataset] Add PPI and PPI-large dataset (THUDM#175)

THINK2TRY · web-flow · commit 952830b9faa6 · 2021-01-25T11:48:58.000+08:00
* Add PPI and PPI-Large

* Fix bugs in PPRGo trainer

* Fix bugs in subgraph sampling

* Simplify node_classification test

* Add tests for graphsage
diff --git a/cogdl/data/data.py b/cogdl/data/data.py
@@ -1,10 +1,22 @@
 import re
+import numba
 
 import torch
 import numpy as np
 import scipy.sparse as sparse
 
 
+@numba.njit
+def reindex(node_idx, col):
+    node_dict = dict()
+    cnt = 0
+    for i in node_idx:
+        node_dict[i] = cnt
+        cnt += 1
+    new_col = [node_dict[i] for i in col]
+    return np.array(new_col)
+
+
 class Data(object):
     r"""A plain old python object modeling a single graph with various
     (optional) attributes:
@@ -184,6 +196,12 @@ def _build_adj_(self):
         if self.__adj is not None:
             return
         num_edges = self.edge_index.shape[1]
+
+        self.edge_row, col = self.edge_index.cpu()
+        self.edge_col = col.numpy()
+
+        self.node_idx = torch.unique(self.edge_index).cpu().numpy()
+
         edge_index_np = self.edge_index.cpu().numpy()
         num_nodes = self.x.shape[0]
         edge_attr_np = np.ones(num_edges)
@@ -205,7 +223,9 @@ def subgraph(self, node_idx):
         edge_attr = torch.from_numpy(adj_coo.data).to(self.x.device)
         edge_index = torch.from_numpy(np.stack([row, col], axis=0)).to(self.x.device).long()
         keys = self.keys
-        attrs = {key: self[key][node_idx] for key in keys if "edge" not in key}
+
+        print(keys)
+        attrs = {key: self[key][node_idx] for key in keys if "edge" not in key and "node_idx" not in key}
         attrs["edge_index"] = edge_index
         if edge_attr is not None:
             attrs["edge_attr"] = edge_attr
@@ -243,37 +263,32 @@ def sample_adj(self, batch, size=-1, replace=True):
         adj = self.__adj[batch].tocsr()
         batch_size = len(batch)
         if size == -1:
-            adj = adj.tocoo()
-            row, col = torch.from_numpy(adj.row), torch.from_numpy(adj.col)
-            node_idx = torch.unique(col)
+            row, col = self.edge_row, self.edge_col
+            _node_idx = self.node_idx
         else:
             indices = torch.from_numpy(adj.indices)
             indptr = torch.from_numpy(adj.indptr)
             node_idx, (row, col) = self._sample_adj(batch_size, indices, indptr, size)
-        col = col.numpy()
-        _node_idx = node_idx.numpy()
+            col = col.numpy()
+            _node_idx = node_idx.numpy()
 
         # Reindexing: target nodes are always put at the front
-        _node_idx = list(batch) + list(set(_node_idx).difference(set(batch)))
-        node_dict = {val: key for key, val in enumerate(_node_idx)}
-        new_col = torch.LongTensor([node_dict[i] for i in col])
+        _node_idx = np.concatenate((batch, np.setdiff1d(_node_idx, batch)))
+
+        new_col = torch.as_tensor(reindex(_node_idx, col), dtype=torch.long)
         edge_index = torch.stack([row.long(), new_col])
 
-        node_idx = torch.Tensor(_node_idx).long().to(self.x.device)
+        node_idx = torch.as_tensor(_node_idx, dtype=torch.long).to(self.x.device)
         edge_index = edge_index.long().to(self.x.device)
         return node_idx, edge_index
 
     def _sample_adj(self, batch_size, indices, indptr, size):
-        indptr = indptr
-        row_counts = torch.Tensor([indptr[i] - indptr[i - 1] for i in range(1, len(indptr))])
+        row_counts = torch.as_tensor([indptr[i] - indptr[i - 1] for i in range(1, len(indptr))]).long()
 
-        # if not replace:
-        #     edge_cols = [col[indptr[i]: indptr[i+1]] for i in range(len(indptr)-1)]
-        #     edge_cols = [np.random.choice(x, min(size, len(x)), replace=False) for x in edge_cols]
-        # else:
         rand = torch.rand(batch_size, size)
         rand = rand * row_counts.view(-1, 1)
         rand = rand.long()
+
         rand = rand + indptr[:-1].view(-1, 1)
         edge_cols = indices[rand].view(-1)
         row = torch.arange(0, batch_size).view(-1, 1).repeat(1, size).view(-1)
diff --git a/cogdl/datasets/__init__.py b/cogdl/datasets/__init__.py
@@ -109,7 +109,7 @@ def build_dataset_from_path(data_path, task):
     "blogcatalog": "cogdl.datasets.matlab_matrix",
     "flickr-ne": "cogdl.datasets.matlab_matrix",
     "wikipedia": "cogdl.datasets.matlab_matrix",
-    "ppi": "cogdl.datasets.matlab_matrix",
+    "ppi-ne": "cogdl.datasets.matlab_matrix",
     "han-acm": "cogdl.datasets.han_data",
     "han-dblp": "cogdl.datasets.han_data",
     "han-imdb": "cogdl.datasets.han_data",
@@ -129,6 +129,8 @@ def build_dataset_from_path(data_path, task):
     "amazon-s": "cogdl.datasets.saint_data",
     "flickr": "cogdl.datasets.saint_data",
     "reddit": "cogdl.datasets.saint_data",
+    "ppi": "cogdl.datasets.saint_data",
+    "ppi-large": "cogdl.datasets.saint_data",
     "test_bio": "cogdl.datasets.pyg_strategies_data",
     "test_chem": "cogdl.datasets.pyg_strategies_data",
     "bio": "cogdl.datasets.strategies_data",
diff --git a/cogdl/datasets/matlab_matrix.py b/cogdl/datasets/matlab_matrix.py
@@ -93,7 +93,7 @@ def __init__(self):
         super(WikipediaDataset, self).__init__(path, filename, url)
 
 
-@register_dataset("ppi")
+@register_dataset("ppi-ne")
 class PPIDataset(MatlabMatrix):
     def __init__(self):
         dataset, filename = "ppi", "Homo_sapiens"
diff --git a/cogdl/datasets/saint_data.py b/cogdl/datasets/saint_data.py
@@ -167,3 +167,33 @@ def __init__(self):
 
     def get_evaluator(self):
         return multiclass_evaluator()
+
+
+@register_dataset("ppi")
+class PPIDataset(SAINTDataset):
+    def __init__(self):
+        dataset = "PPI"
+        url = "https://cloud.tsinghua.edu.cn/d/2c6e94ec9dad4972b58e/files/?p=%2F{}&dl=1"
+        path = osp.join("data", dataset)
+        if not osp.exists(path):
+            SAINTDataset(path, dataset, url)
+        super(PPIDataset, self).__init__(path, dataset, url)
+        self.data = scale_feats(self.data)
+
+    def get_evaluator(self):
+        return multilabel_evaluator()
+
+
+@register_dataset("ppi-large")
+class PPILargeDataset(SAINTDataset):
+    def __init__(self):
+        dataset = "PPI_Large"
+        url = "https://cloud.tsinghua.edu.cn/d/469ef38a520640bba267/files/?p=%2F{}&dl=1"
+        path = osp.join("data", dataset)
+        if not osp.exists(path):
+            SAINTDataset(path, dataset, url)
+        super(PPILargeDataset, self).__init__(path, dataset, url)
+        self.data = scale_feats(self.data)
+
+    def get_evaluator(self):
+        return multilabel_evaluator()
diff --git a/cogdl/match.yml b/cogdl/match.yml
@@ -33,8 +33,9 @@ node_classification:
     - ogbn-products
     - ogbn-papers100M
     - flickr
-    - amazon
+    - amazon-s
     - yelp
+    - ppi
     - reddit
 unsupervised_node_classification:
   - model:
@@ -51,7 +52,7 @@ unsupervised_node_classification:
     - spectral
     - gcc
     dataset:
-    - ppi
+    - ppi-ne
     - blogcatalog
     - wikipedia
     - usa-airport
diff --git a/cogdl/models/nn/graphsage.py b/cogdl/models/nn/graphsage.py
@@ -159,6 +159,8 @@ def inference(self, x_all, data_loader):
     def get_trainer(task: Any, args: Any):
         if args.dataset not in ["cora", "citeseer"]:
             return NeighborSamplingTrainer
+        if hasattr(args, "use_trainer"):
+            return NeighborSamplingTrainer
 
     def set_data_device(self, device):
         self.device = device
diff --git a/cogdl/models/nn/pyg_graph_unet.py b/cogdl/models/nn/pyg_graph_unet.py
@@ -182,6 +182,7 @@ def __init__(
         self.act = get_activation(activation)
         assert pooling_layer <= len(pooling_rates)
         pooling_rates = pooling_rates[:pooling_layer]
+        pooling_rates = [float(x) for x in pooling_rates]
         self.unet = GraphUnetLayer(hidden_size, pooling_layer, pooling_rates, activation, n_dropout, aug_adj)
 
         self.in_gcn = GraphConvolution(in_feats, hidden_size)
diff --git a/cogdl/trainers/ppr_trainer.py b/cogdl/trainers/ppr_trainer.py
@@ -50,9 +50,9 @@ def ppr_run(self, dataset, mode="train"):
         else:
             edge_index = data.edge_index
 
-        if not os.path.exists("./saved"):
-            os.mkdir("saved")
-        path = f"./saved/{self.dataset_name}_{self.topk}_{self.alpha}_{self.normalization}.{mode}.npz"
+        if not os.path.exists("./pprgo_saved"):
+            os.mkdir("pprgo_saved")
+        path = f"./pprgo_saved/{self.dataset_name}_{self.topk}_{self.alpha}_{self.normalization}.{mode}.npz"
 
         if os.path.exists(path):
             print(f"Load {mode} from cached")
@@ -103,7 +103,6 @@ def fit(self, model, dataset):
             epoch_iter.set_description(
                 f"Epoch: {epoch}, TrainLoss: {train_loss: .4f}, ValLoss: {val_loss: .4f}, ValAcc: {best_acc: .4f}"
             )
-            print()
         self.model = best_model
 
         del train_loader
diff --git a/cogdl/trainers/sampled_trainer.py b/cogdl/trainers/sampled_trainer.py
@@ -7,10 +7,11 @@
 import torch.nn.functional as F
 from tqdm import tqdm
 
-from cogdl.data import Dataset, Data
+from cogdl.data import Dataset
 from cogdl.data.sampler import NodeSampler, EdgeSampler, RWSampler, MRWSampler, LayerSampler, NeighborSampler
 from cogdl.models.supervised_model import SupervisedModel
 from cogdl.trainers.base_trainer import BaseTrainer
+from cogdl.utils import add_remaining_self_loops
 
 
 class SampledTrainer(BaseTrainer):
@@ -149,9 +150,14 @@ def __init__(self, args):
         self.patience = self.patience // self.eval_per_epoch
 
         self.device = "cpu" if not torch.cuda.is_available() or args.cpu else args.device_id[0]
+        self.loss_fn, self.evaluator = None, None
 
     def fit(self, model, dataset):
-        self.data = Data.from_pyg_data(dataset[0])
+        self.data = dataset[0]
+        self.data.edge_index, _ = add_remaining_self_loops(self.data.edge_index)
+        if hasattr(self.data, "edge_index_train"):
+            self.data.edge_index_train, _ = add_remaining_self_loops(self.data.edge_index_train)
+        self.loss_fn, self.evaluator = dataset.get_evaluator()
         self.train_loader = NeighborSampler(
             data=self.data,
             mask=self.data.train_mask,
diff --git a/scripts/train.py b/scripts/train.py
@@ -5,7 +5,6 @@
     parser = options.get_training_parser()
     args, _ = parser.parse_known_args()
     args = options.parse_args_and_arch(parser, args)
-    print(args)
     assert len(args.device_id) == 1
 
     experiment(task=args.task, dataset=args.dataset, model=args.model, args=args)
diff --git a/tests/datasets/test_data.py b/tests/datasets/test_data.py
@@ -1,13 +1,12 @@
 import numpy as np
 
-from cogdl.data import Data
 from cogdl.datasets import build_dataset_from_name
 
 
 class Test_Data(object):
     def setup_class(self):
         self.dataset = build_dataset_from_name("cora")
-        self.data = Data.from_pyg_data(self.dataset[0])
+        self.data = self.dataset[0]
         self.num_nodes = self.data.num_nodes
         self.num_edges = self.data.num_edges
         self.num_features = self.data.num_features
@@ -33,4 +32,3 @@ def test_adj_sampling(self):
             node_idx, sampled_edge_index = self.data.sample_adj(sampled_nodes, size)
             node_idx = node_idx.cpu().numpy()
             assert (set(node_idx) & set(sampled_nodes)) == set(sampled_nodes)
-            assert len(set(sampled_edge_index.cpu().numpy()[0])) == len(set(sampled_nodes))
diff --git a/tests/datasets/test_saint_data.py b/tests/datasets/test_saint_data.py
@@ -3,6 +3,7 @@
 import numpy as np
 import scipy.sparse as sp
 from cogdl.datasets.saint_data import SAINTDataset
+from cogdl.datasets import build_dataset_from_name
 
 
 def build_test_data():
diff --git a/tests/tasks/test_node_classification.py b/tests/tasks/test_node_classification.py
@@ -190,6 +190,11 @@ def test_graphsage_cora():
     ret = task.train()
     assert 0 <= ret["Acc"] <= 1
 
+    args.use_trainer = True
+    task = build_task(args)
+    ret = task.train()
+    assert 0 <= ret["Acc"] <= 1
+
 
 def test_pyg_cheb_cora():
     args = get_default_args()
@@ -236,10 +241,11 @@ def test_pyg_gcn_cora_sampler():
         assert 0 <= ret["Acc"] <= 1
 
 
-def test_unet_cora():
+def test_unet_citeseer():
     args = get_default_args()
     args.cpu = True
     args.model = "unet"
+    args.dataset = "citeseer"
     args.pool_rate = [0.5, 0.5]
     args.n_pool = 2
     args.adj_dropout = 0.3
@@ -248,6 +254,7 @@ def test_unet_cora():
     args.improved = True
     args.aug_adj = True
     args.activation = "elu"
+    # print(args)
     task = build_task(args)
     ret = task.train()
     assert 0 <= ret["Acc"] <= 1
@@ -662,5 +669,3 @@ def test_pprgo_cora():
     test_dropedge_resgcn_cora()
     test_dropedge_inceptiongcn_cora()
     test_dropedge_densegcn_cora()
-    test_unet_cora()
-    test_pprgo_cora()
diff --git a/tests/tasks/test_node_classification_sampling.py b/tests/tasks/test_node_classification_sampling.py
@@ -1,8 +1,5 @@
 import torch
-from cogdl import options
 from cogdl.tasks import build_task
-from cogdl.datasets import build_dataset
-from cogdl.models import build_model
 from cogdl.utils import build_args_from_dict
 
 
diff --git a/tests/tasks/test_unsupervised_node_classification.py b/tests/tasks/test_unsupervised_node_classification.py
