[Example] Add clustergcn for ogbn (#320)

* Add clustergcn example

* Fix tests

* Fix num_nodes & oagbert

* Fix agc

* Fix tests

* Fix spmm(cpu)

README.md

@@ -14,10 +14,9 @@ CogDL is a graph deep learning toolkit that allows researchers and developers to
 
 We summarize the contributions of CogDL as follows:
 
-- **High Efficiency**: CogDL utilizes well-optimized operators to speed up training and save GPU memory of GNN models.
-- **Easy-to-Use**: CogDL provides easy-to-use APIs for running experiments with the given models and datasets using hyper-parameter search.
+- **Efficiency**: CogDL utilizes well-optimized operators to speed up training and save GPU memory of GNN models.
+- **Ease of Use**: CogDL provides easy-to-use APIs for running experiments with the given models and datasets using hyper-parameter search.
 - **Extensibility**: The design of CogDL makes it easy to apply GNN models to new scenarios based on our framework.
-- **Reproducibility**: CogDL provides reproducible leaderboards for state-of-the-art models on most of important tasks in the graph domain.
 
 ## ❗ News
 
@@ -223,7 +222,7 @@ Please cite [our paper](https://arxiv.org/abs/2103.00959) if you find our code o
 
 ```
 @article{cen2021cogdl,
-    title={CogDL: An Extensive Toolkit for Deep Learning on Graphs},
+    title={CogDL: Toolkit for Deep Learning on Graphs},
     author={Yukuo Cen and Zhenyu Hou and Yan Wang and Qibin Chen and Yizhen Luo and Xingcheng Yao and Aohan Zeng and Shiguang Guo and Peng Zhang and Guohao Dai and Yu Wang and Chang Zhou and Hongxia Yang and Jie Tang},
     journal={arXiv preprint arXiv:2103.00959},
     year={2021}

README_CN.md

@@ -14,10 +14,9 @@ CogDL是一款图深度学习工具包，基于[PyTorch](https://github.com/pyto
 
 CogDL的特性包括：
 
-- 高效：CogDL支持使用优化好的算子来加速GNN模型的训练。
+- 高效性：CogDL支持使用优化好的算子来加速GNN模型的训练。
 - 易用性：CogDL提供了非常易用的API来在给定的模型和数据集上运行实验。
-- 可扩展性：用户可以基于CogDL已有的框架来实现和提交新的数据集、模型和任务。
-- 可复现性：CogDL对图领域大多数重要的任务都提供了可复现的排行榜。
+- 扩展性：用户可以基于CogDL已有的框架来扩展新的数据集、模型。
 
 ## ❗ 最新
 
@@ -200,7 +199,7 @@ CogDL核心开发团队可以通过[cogdlteam@gmail.com](mailto:cogdlteam@gmail.
 
 ```
 @article{cen2021cogdl,
-    title={CogDL: An Extensive Toolkit for Deep Learning on Graphs},
+    title={CogDL: Toolkit for Deep Learning on Graphs},
     author={Yukuo Cen and Zhenyu Hou and Yan Wang and Qibin Chen and Yizhen Luo and Xingcheng Yao and Aohan Zeng and Shiguang Guo and Peng Zhang and Guohao Dai and Yu Wang and Chang Zhou and Hongxia Yang and Jie Tang},
     journal={arXiv preprint arXiv:2103.00959},
     year={2021}

cogdl/__init__.py

@@ -1,5 +1,4 @@
 __version__ = "0.5.1.post1"
 
 from .experiments import experiment
-from .oag import oagbert
 from .pipelines import pipeline

cogdl/data/data.py

@@ -17,6 +17,8 @@
 from cogdl.utils import RandomWalker
 from cogdl.operators.sample import sample_adj_c, subgraph_c
 
+subgraph_c = None  # noqa: F811
+
 
 class BaseGraph(object):
     def __init__(self):
@@ -628,7 +630,11 @@ def edge_index(self, edge_index):
                 self._adj.row_ptr = None
             self._adj.row = row
             self._adj.col = col
-            self.__num_nodes__ = None
+            if self.x is not None:
+                self._adj.__num_nodes__ = self.x.shape[0]
+                self.__num_nodes__ = self.x.shape[0]
+            else:
+                self.__num_nodes__ = None
 
     @edge_weight.setter
     def edge_weight(self, edge_weight):
@@ -857,6 +863,7 @@ def csr_subgraph(self, node_idx, keep_order=False):
                 continue
             data._adj[key] = self._adj[key][edges]
         data.num_nodes = node_idx.shape[0]
+        data.edge_weight = None
         return data
 
     def subgraph(self, node_idx, keep_order=False):

cogdl/layers/sage_layer.py

@@ -51,11 +51,6 @@ def __init__(
         else:
             raise NotImplementedError
 
-        if dropout > 0:
-            self.dropout = nn.Dropout(dropout)
-        else:
-            self.dropout = None
-
         if activation is not None:
             self.act = get_activation(activation, inplace=True)
         else:

cogdl/models/nn/daegc.py

@@ -81,4 +81,4 @@ def get_features(self, data):
 
     def recon_loss(self, z, adj):
         # print(torch.mm(z, z.t()), adj)
-        return F.binary_cross_entropy(F.softmax(torch.mm(z, z.t())), adj, reduction="sum")
+        return F.binary_cross_entropy(F.sigmoid(torch.mm(z, z.t())), adj, reduction="sum")

cogdl/models/nn/gae.py

@@ -47,6 +47,7 @@ def __init__(self, num_features, hidden_size):
         self.conv2_var = GCNLayer(self.hidden_size, self.hidden_size)
 
     def reparameterize(self, mean, log_var):
+        log_var = log_var.clamp(max=10)
         sigma = torch.exp(log_var)
         z = mean + torch.randn_like(log_var) * sigma
         return z

cogdl/models/nn/infograph.py

@@ -88,7 +88,7 @@ def forward(self, x):
 
 
 class InfoGraph(BaseModel):
-    r"""Implimentation of Infograph in paper `"InfoGraph: Unsupervised and Semi-supervised Graph-Level Representation
+    r"""Implementation of Infograph in paper `"InfoGraph: Unsupervised and Semi-supervised Graph-Level Representation
      Learning via Mutual Information Maximization" <https://openreview.net/forum?id=r1lfF2NYvH>__. `
 
      Parameters

cogdl/oag/README.md

@@ -8,7 +8,7 @@ A basic version OAG-BERT. Similar to [SciBERT](https://github.com/allenai/sciber
 
 The usage of OAG-BERT is the same of ordinary SciBERT or BERT. For example, you can use the following code to encode two text sequences and retrieve their outputs
 ```python
-from cogdl import oagbert
+from cogdl.oag import oagbert
 
 tokenizer, bert_model = oagbert()
 
@@ -20,7 +20,7 @@ outputs = bert_model(**tokens)
 ## V2: The entity augmented version
 An extension to the vanilla OAG-BERT. We incorporate rich entity information in Open Academic Graph such as **authors** and **field-of-study**. Thus, you can encode various type of entities in OAG-BERT v2. For example, to encode the paper of BERT, you can use the following code
 ```python
-from cogdl import oagbert
+from cogdl.oag import oagbert
 import torch
 
 tokenizer, model = oagbert("oagbert-v2")
@@ -47,7 +47,7 @@ sequence_output, pooled_output = model.bert.forward(
 ```
 If you want to encode various type of entities separately, you can use the following code instead
 ```python
-from cogdl import oagbert
+from cogdl.oag import oagbert
 
 tokenizer, model = oagbert("oagbert-v2")
 title = 'BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding'
@@ -87,7 +87,7 @@ We also release another two V2 version for users.
 
 One is a generation based version which can be used for generating texts based on other information. For example, use the following code to automatically generate paper titles with abstracts.
 ```python
-from cogdl import oagbert
+from cogdl.oag import oagbert
 
 tokenizer, model = oagbert('oagbert-v2-lm')
 model.eval()
@@ -101,7 +101,7 @@ for seq, prob in model.generate_title(abstract="To enrich language models with d
 In addition to that, we fine-tune the OAG-BERT for calculating paper similarity based on name disambiguation tasks, which is named as Sentence-OAGBERT following [Sentence-BERT](https://github.com/UKPLab/sentence-transformers). The following codes demonstrate an example of using Sentence-OAGBERT to calculate paper similarity.
 ```python
 import os
-from cogdl import oagbert
+from cogdl.oag import oagbert
 import torch
 import torch.nn.functional as F
 import numpy as np

cogdl/operators/spmm/spmm_cpu.cpp

@@ -13,7 +13,7 @@ torch::Tensor spmm_cpu(
     const auto k = dense.size(1);
     auto devid = dense.device().index();
     auto options = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCPU, devid);
-    auto out = torch::empty({m,k}, options);
+    auto out = torch::zeros({m,k}, options);
 
     int *rowptr_ptr = rowptr.data_ptr<int>();
     int *colind_ptr = colind.data_ptr<int>();

cogdl/pipelines.py

@@ -9,7 +9,6 @@
 from grave import plot_network, use_attributes
 from tabulate import tabulate
 
-from cogdl import oagbert
 from cogdl.data import Graph
 from cogdl.datasets import build_dataset_from_name, NodeDataset
 from cogdl.models import build_model
@@ -126,6 +125,9 @@ def __init__(self, app: str, model: str, **kwargs):
         super(OAGBertInferencePipepline, self).__init__(app, model=model, **kwargs)
 
         load_weights = kwargs["load_weights"] if "load_weights" in kwargs else True
+
+        from cogdl.oag import oagbert
+
         self.tokenizer, self.bert_model = oagbert(model, load_weights=load_weights)
 
     def __call__(self, sequence, **kwargs):

cogdl/trainer/trainer.py

@@ -275,7 +275,7 @@ def initialize(self, model_w, rank=0, master_addr: str = "localhost", master_por
         else:
             return model_w.to(rank), None
 
-    def train(self, rank, model_w, dataset_w):
+    def train(self, rank, model_w, dataset_w):  # noqa: C901
         model_w, _ = self.initialize(model_w, rank=rank, master_addr=self.master_addr, master_port=self.master_port)
         self.data_controller.prepare_data_wrapper(dataset_w, rank)
         self.eval_data_back_to_cpu = dataset_w.data_back_to_cpu
@@ -306,10 +306,10 @@ def train(self, rank, model_w, dataset_w):
                 self.data_controller.training_proc_per_stage(dataset_w, rank)
 
             if self.progress_bar == "epoch":
-                epoch_iter = tqdm(range(self.epochs))
+                epoch_iter = tqdm(range(1, self.epochs + 1))
                 epoch_printer = Printer(epoch_iter.set_description, rank=rank, world_size=self.world_size)
             else:
-                epoch_iter = range(self.epochs)
+                epoch_iter = range(1, self.epochs + 1)
                 epoch_printer = Printer(print, rank=rank, world_size=self.world_size)
 
             self.logger.start()
@@ -321,7 +321,10 @@ def train(self, rank, model_w, dataset_w):
                 # inductive setting ..
                 dataset_w.train()
                 train_loader = dataset_w.on_train_wrapper()
-                training_loss = self.training_step(model_w, train_loader, optimizers, lr_schedulers, rank)
+                train_dataset = train_loader.get_dataset_from_loader()
+                if hasattr(train_dataset, "shuffle"):
+                    train_dataset.shuffle()
+                training_loss = self.train_step(model_w, train_loader, optimizers, lr_schedulers, rank)
 
                 print_str_dict["Epoch"] = epoch
                 print_str_dict["train_loss"] = training_loss
@@ -386,12 +389,13 @@ def validate(self, model_w: ModelWrapper, dataset_w: DataWrapper, device):
         dataset_w.eval()
         if self.cpu_inference:
             model_w.to("cpu")
-            _device = device
+            _device = "cpu"
         else:
             _device = device
 
         val_loader = dataset_w.on_val_wrapper()
-        result = self.val_step(model_w, val_loader, _device)
+        with torch.no_grad():
+            result = self.val_step(model_w, val_loader, _device)
 
         model_w.to(device)
         return result
@@ -406,12 +410,16 @@ def test(self, model_w: ModelWrapper, dataset_w: DataWrapper, device):
         dataset_w.eval()
         if self.cpu_inference:
             model_w.to("cpu")
-            _device = device
+            _device = "cpu"
         else:
             _device = device
 
         test_loader = dataset_w.on_test_wrapper()
-        result = self.test_step(model_w, test_loader, _device)
+        if model_w.training_type == "unsupervised":
+            result = self.test_step(model_w, test_loader, _device)
+        else:
+            with torch.no_grad():
+                result = self.test_step(model_w, test_loader, _device)
 
         model_w.to(device)
         return result
@@ -425,7 +433,8 @@ def distributed_test(self, model_w: ModelWrapper, loader, rank, fn):
                 _device = "cpu"
             else:
                 _device = rank
-            result = fn(model_w, loader, _device)
+            with torch.no_grad():
+                result = fn(model_w, loader, _device)
             model_w.to(rank)
 
             object_list = [result]
@@ -434,7 +443,7 @@ def distributed_test(self, model_w: ModelWrapper, loader, rank, fn):
         dist.broadcast_object_list(object_list, src=0)
         return object_list[0]
 
-    def training_step(self, model_w, train_loader, optimizers, lr_schedulers, device):
+    def train_step(self, model_w, train_loader, optimizers, lr_schedulers, device):
         model_w.train()
         losses = []
 
@@ -443,6 +452,8 @@ def training_step(self, model_w, train_loader, optimizers, lr_schedulers, device
 
         for batch in train_loader:
             batch = move_to_device(batch, device)
+            if hasattr(batch, "train_mask") and batch.train_mask.sum().item() == 0:
+                continue
             loss = model_w.on_train_step(batch)
 
             for optimizer in optimizers:
@@ -458,9 +469,9 @@ def training_step(self, model_w, train_loader, optimizers, lr_schedulers, device
         if lr_schedulers is not None:
             for lr_schedular in lr_schedulers:
                 lr_schedular.step()
+
         return np.mean(losses)
 
-    @torch.no_grad()
     def val_step(self, model_w, val_loader, device):
         model_w.eval()
         if val_loader is None:
@@ -472,7 +483,6 @@ def val_step(self, model_w, val_loader, device):
                 move_to_device(batch, "cpu")
         return model_w.collect_notes()
 
-    # @torch.no_grad()
     def test_step(self, model_w, test_loader, device):
         model_w.eval()
         if test_loader is None:

cogdl/utils/spmm_utils.py

@@ -82,10 +82,13 @@ def forward(self, graph, x):
         return spmm_cpu(graph, x, self.fast_spmm_cpu)
 
 
-def spmm(graph, x, actnn=False, fast_spmm=None):
+def spmm(graph, x, actnn=False, fast_spmm=None, fast_spmm_cpu=None):
     if fast_spmm is None:
         initialize_spmm()
         fast_spmm = CONFIGS["fast_spmm"]
+    if fast_spmm_cpu is None:
+        initialize_spmm_cpu()
+        fast_spmm_cpu = CONFIGS["fast_spmm_cpu"]
     if fast_spmm is not None and str(x.device) != "cpu":
         if graph.out_norm is not None:
             x = graph.out_norm * x
@@ -94,6 +97,16 @@ def spmm(graph, x, actnn=False, fast_spmm=None):
         csr_data = graph.raw_edge_weight
         x = fast_spmm(row_ptr.int(), col_indices.int(), x, csr_data, graph.is_symmetric(), actnn=actnn)
 
+        if graph.in_norm is not None:
+            x = graph.in_norm * x
+    elif fast_spmm_cpu is not None and str(x.device) == "cpu" and x.requires_grad is False:
+        if graph.out_norm is not None:
+            x = graph.out_norm * x
+
+        row_ptr, col_indices = graph.row_indptr, graph.col_indices
+        csr_data = graph.raw_edge_weight
+        x = fast_spmm_cpu(row_ptr.int(), col_indices.int(), csr_data, x)
+
         if graph.in_norm is not None:
             x = graph.in_norm * x
     else:

cogdl/wrappers/model_wrapper/__init__.py

@@ -1,6 +1,6 @@
 import importlib
 
-from .base_model_wrapper import ModelWrapper, EmbeddingModelWrapper
+from .base_model_wrapper import ModelWrapper, EmbeddingModelWrapper, UnsupervisedModelWrapper
 
 
 def register_model_wrapper(name):

cogdl/wrappers/model_wrapper/base_model_wrapper.py

@@ -17,6 +17,7 @@ def __init__(self):
         self._evaluator = None
         self._evaluator_metric = None
         self.__record__ = dict()
+        self.training_type = ""
 
     def forward(self):
         pass
@@ -196,3 +197,9 @@ def _model_key_(self):
 class EmbeddingModelWrapper(ModelWrapper):
     def setup_optimizer(self):
         pass
+
+
+class UnsupervisedModelWrapper(ModelWrapper):
+    def __init__(self):
+        super(UnsupervisedModelWrapper, self).__init__()
+        self.training_type = "unsupervised"

cogdl/wrappers/model_wrapper/node_classification/dgi_mw.py

@@ -3,11 +3,11 @@
 import torch
 import torch.nn as nn
 
-from .. import ModelWrapper
+from .. import UnsupervisedModelWrapper
 from cogdl.wrappers.tools.wrapper_utils import evaluate_node_embeddings_using_logreg
 
 
-class DGIModelWrapper(ModelWrapper):
+class DGIModelWrapper(UnsupervisedModelWrapper):
     @staticmethod
     def add_args(parser):
         # fmt: off

cogdl/wrappers/model_wrapper/node_classification/grace_mw.py

@@ -3,12 +3,12 @@
 import torch.nn.functional as F
 
 from cogdl.data import Graph
-from .. import ModelWrapper
+from .. import UnsupervisedModelWrapper
 from cogdl.wrappers.tools.wrapper_utils import evaluate_node_embeddings_using_logreg
 from cogdl.utils import dropout_adj, dropout_features
 
 
-class GRACEModelWrapper(ModelWrapper):
+class GRACEModelWrapper(UnsupervisedModelWrapper):
     @staticmethod
     def add_args(parser):
         # fmt: off