chore: import upstream snapshot with attribution

python/dgl/data/__init__.py

@@ -0,0 +1,112 @@
+"""The ``dgl.data`` package contains datasets hosted by DGL and also utilities
+for downloading, processing, saving and loading data from external resources.
+"""
+
+from __future__ import absolute_import
+
+from . import citation_graph as citegrh
+from .actor import ActorDataset
+from .movielens import MovieLensDataset
+from .adapter import *
+from .bitcoinotc import BitcoinOTC, BitcoinOTCDataset
+from .citation_graph import (
+    CitationGraphDataset,
+    CiteseerGraphDataset,
+    CoraBinary,
+    CoraGraphDataset,
+    PubmedGraphDataset,
+)
+from .csv_dataset import CSVDataset
+from .dgl_dataset import DGLBuiltinDataset, DGLDataset
+from .fakenews import FakeNewsDataset
+from .flickr import FlickrDataset
+from .fraud import FraudAmazonDataset, FraudDataset, FraudYelpDataset
+from .gdelt import GDELT, GDELTDataset
+from .gindt import GINDataset
+from .gnn_benchmark import (
+    AmazonCoBuy,
+    AmazonCoBuyComputerDataset,
+    AmazonCoBuyPhotoDataset,
+    Coauthor,
+    CoauthorCSDataset,
+    CoauthorPhysicsDataset,
+    CoraFull,
+    CoraFullDataset,
+)
+from .icews18 import ICEWS18, ICEWS18Dataset
+from .karate import KarateClub, KarateClubDataset
+from .knowledge_graph import FB15k237Dataset, FB15kDataset, WN18Dataset
+from .minigc import *
+from .ppi import LegacyPPIDataset, PPIDataset
+from .qm7b import QM7b, QM7bDataset
+from .qm9 import QM9, QM9Dataset
+from .qm9_edge import QM9Edge, QM9EdgeDataset
+from .rdf import AIFBDataset, AMDataset, BGSDataset, MUTAGDataset
+from .reddit import RedditDataset
+from .sbm import SBMMixture, SBMMixtureDataset
+from .synthetic import (
+    BA2MotifDataset,
+    BACommunityDataset,
+    BAShapeDataset,
+    TreeCycleDataset,
+    TreeGridDataset,
+)
+from .tree import SST, SSTDataset
+from .tu import LegacyTUDataset, TUDataset
+from .utils import *
+from .cluster import CLUSTERDataset
+from .geom_gcn import (
+    ChameleonDataset,
+    CornellDataset,
+    SquirrelDataset,
+    TexasDataset,
+    WisconsinDataset,
+)
+
+from .heterophilous_graphs import (
+    AmazonRatingsDataset,
+    MinesweeperDataset,
+    QuestionsDataset,
+    RomanEmpireDataset,
+    TolokersDataset,
+)
+
+# RDKit is required for Peptides-Structural, Peptides-Functional dataset.
+# Exception handling was added to prevent crashes for users who are using other
+# datasets.
+try:
+    from .lrgb import (
+        COCOSuperpixelsDataset,
+        PeptidesFunctionalDataset,
+        PeptidesStructuralDataset,
+        VOCSuperpixelsDataset,
+    )
+except ImportError:
+    pass
+from .pattern import PATTERNDataset
+from .superpixel import CIFAR10SuperPixelDataset, MNISTSuperPixelDataset
+from .wikics import WikiCSDataset
+from .yelp import YelpDataset
+from .zinc import ZINCDataset
+
+
+def register_data_args(parser):
+    parser.add_argument(
+        "--dataset",
+        type=str,
+        required=False,
+        help="The input dataset. Can be cora, citeseer, pubmed, syn(synthetic dataset) or reddit",
+    )
+
+
+def load_data(args):
+    if args.dataset == "cora":
+        return citegrh.load_cora()
+    elif args.dataset == "citeseer":
+        return citegrh.load_citeseer()
+    elif args.dataset == "pubmed":
+        return citegrh.load_pubmed()
+    elif args.dataset is not None and args.dataset.startswith("reddit"):
+        return RedditDataset(self_loop=("self-loop" in args.dataset))
+    else:
+        raise ValueError("Unknown dataset: {}".format(args.dataset))

python/dgl/data/actor.py

@@ -0,0 +1,138 @@
+"""
+Actor-only induced subgraph of the film-directoractor-writer network.
+"""
+import os
+
+import numpy as np
+
+from ..convert import graph
+from .dgl_dataset import DGLBuiltinDataset
+from .utils import _get_dgl_url
+
+
+class ActorDataset(DGLBuiltinDataset):
+    r"""Actor-only induced subgraph of the film-directoractor-writer network
+    from `Social Influence Analysis in Large-scale Networks
+    <https://dl.acm.org/doi/10.1145/1557019.1557108>`, introduced by
+    `Geom-GCN: Geometric Graph Convolutional Networks
+    <https://arxiv.org/abs/2002.05287>`
+
+    Nodes represent actors, and edges represent co-occurrence on the same
+    Wikipedia page. Node features correspond to some keywords in the Wikipedia
+    pages.
+
+    Statistics:
+
+    - Nodes: 7600
+    - Edges: 33391
+    - Number of Classes: 5
+    - 10 train/val/test splits
+
+        - Train: 3648
+        - Val: 2432
+        - Test: 1520
+
+    Parameters
+    ----------
+    raw_dir : str, optional
+        Raw file directory to store the processed data. Default: ~/.dgl/
+    force_reload : bool, optional
+        Whether to re-download the data source. Default: False
+    verbose : bool, optional
+        Whether to print progress information. Default: True
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access. Default: None
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of node classes
+
+    Notes
+    -----
+    The graph does not come with edges for both directions.
+    """
+
+    def __init__(
+        self, raw_dir=None, force_reload=False, verbose=True, transform=None
+    ):
+        super(ActorDataset, self).__init__(
+            name="actor",
+            url=_get_dgl_url("dataset/actor.zip"),
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        """Load and process the data."""
+        try:
+            import torch
+        except ImportError:
+            raise ModuleNotFoundError(
+                "This dataset requires PyTorch to be the backend."
+            )
+
+        # Process node features and labels.
+        with open(f"{self.raw_path}/out1_node_feature_label.txt", "r") as f:
+            data = [x.split("\t") for x in f.read().split("\n")[1:-1]]
+
+            rows, cols = [], []
+            labels = torch.empty(len(data), dtype=torch.long)
+            for n_id, col, label in data:
+                col = [int(x) for x in col.split(",")]
+                rows += [int(n_id)] * len(col)
+                cols += col
+
+                labels[int(n_id)] = int(label)
+
+            row, col = torch.tensor(rows), torch.tensor(cols)
+            features = torch.zeros(len(data), int(col.max()) + 1)
+            features[row, col] = 1.0
+
+            self._num_classes = int(labels.max().item()) + 1
+
+        # Process graph structure.
+        with open(f"{self.raw_path}/out1_graph_edges.txt", "r") as f:
+            data = f.read().split("\n")[1:-1]
+            data = [[int(v) for v in r.split("\t")] for r in data]
+        dst, src = torch.tensor(data, dtype=torch.long).t().contiguous()
+
+        self._g = graph((src, dst), num_nodes=features.size(0))
+        self._g.ndata["feat"] = features
+        self._g.ndata["label"] = labels
+
+        # Process 10 train/val/test node splits.
+        train_masks, val_masks, test_masks = [], [], []
+        for i in range(10):
+            filepath = f"{self.raw_path}/{self.name}_split_0.6_0.2_{i}.npz"
+            f = np.load(filepath)
+            train_masks += [torch.from_numpy(f["train_mask"])]
+            val_masks += [torch.from_numpy(f["val_mask"])]
+            test_masks += [torch.from_numpy(f["test_mask"])]
+        self._g.ndata["train_mask"] = torch.stack(train_masks, dim=1).bool()
+        self._g.ndata["val_mask"] = torch.stack(val_masks, dim=1).bool()
+        self._g.ndata["test_mask"] = torch.stack(test_masks, dim=1).bool()
+
+    def has_cache(self):
+        return os.path.exists(self.raw_path)
+
+    def load(self):
+        self.process()
+
+    def __getitem__(self, idx):
+        assert idx == 0, "This dataset has only one graph."
+        if self._transform is None:
+            return self._g
+        else:
+            return self._transform(self._g)
+
+    def __len__(self):
+        return 1
+
+    @property
+    def num_classes(self):
+        return self._num_classes

python/dgl/data/adapter.py

@@ -0,0 +1,633 @@
+"""Dataset adapters for re-purposing a dataset for a different kind of training task."""
+
+import json
+import os
+
+import numpy as np
+
+from .. import backend as F
+from ..base import DGLError
+from ..convert import graph as create_dgl_graph
+from ..sampling.negative import _calc_redundancy
+from . import utils
+from .dgl_dataset import DGLDataset
+
+__all__ = ["AsNodePredDataset", "AsLinkPredDataset", "AsGraphPredDataset"]
+
+
+class AsNodePredDataset(DGLDataset):
+    """Repurpose a dataset for a standard semi-supervised transductive
+    node prediction task.
+
+    The class converts a given dataset into a new dataset object such that:
+
+      - Contains only one graph, accessible from ``dataset[0]``.
+      - The graph stores:
+
+        - Node labels in ``g.ndata['label']``.
+        - Train/val/test masks in ``g.ndata['train_mask']``, ``g.ndata['val_mask']``,
+          and ``g.ndata['test_mask']`` respectively.
+      - In addition, the dataset contains the following attributes:
+
+        - ``num_classes``, the number of classes to predict.
+        - ``train_idx``, ``val_idx``, ``test_idx``, train/val/test indexes.
+
+    If the input dataset contains heterogeneous graphs, users need to specify the
+    ``target_ntype`` argument to indicate which node type to make predictions for.
+    In this case:
+
+      - Node labels are stored in ``g.nodes[target_ntype].data['label']``.
+      - Training masks are stored in ``g.nodes[target_ntype].data['train_mask']``.
+        So do validation and test masks.
+
+    The class will keep only the first graph in the provided dataset and
+    generate train/val/test masks according to the given split ratio. The generated
+    masks will be cached to disk for fast re-loading. If the provided split ratio
+    differs from the cached one, it will re-process the dataset properly.
+
+    Parameters
+    ----------
+    dataset : DGLDataset
+        The dataset to be converted.
+    split_ratio : (float, float, float), optional
+        Split ratios for training, validation and test sets. They must sum to one.
+    target_ntype : str, optional
+        The node type to add split mask for.
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of classes to predict.
+    train_idx : Tensor
+        An 1-D integer tensor of training node IDs.
+    val_idx : Tensor
+        An 1-D integer tensor of validation node IDs.
+    test_idx : Tensor
+        An 1-D integer tensor of test node IDs.
+
+    Examples
+    --------
+    >>> ds = dgl.data.AmazonCoBuyComputerDataset()
+    >>> print(ds)
+    Dataset("amazon_co_buy_computer", num_graphs=1, save_path=...)
+    >>> new_ds = dgl.data.AsNodePredDataset(ds, [0.8, 0.1, 0.1])
+    >>> print(new_ds)
+    Dataset("amazon_co_buy_computer-as-nodepred", num_graphs=1, save_path=...)
+    >>> print('train_mask' in new_ds[0].ndata)
+    True
+    """
+
+    def __init__(self, dataset, split_ratio=None, target_ntype=None, **kwargs):
+        self.dataset = dataset
+        self.split_ratio = split_ratio
+        self.target_ntype = target_ntype
+        super().__init__(
+            self.dataset.name + "-as-nodepred",
+            hash_key=(split_ratio, target_ntype, dataset.name, "nodepred"),
+            **kwargs
+        )
+
+    def process(self):
+        is_ogb = hasattr(self.dataset, "get_idx_split")
+        if is_ogb:
+            g, label = self.dataset[0]
+            self.g = g.clone()
+            self.g.ndata["label"] = F.reshape(label, (g.num_nodes(),))
+        else:
+            self.g = self.dataset[0].clone()
+
+        if "label" not in self.g.nodes[self.target_ntype].data:
+            raise ValueError(
+                "Missing node labels. Make sure labels are stored "
+                "under name 'label'."
+            )
+
+        if self.split_ratio is None:
+            if is_ogb:
+                split = self.dataset.get_idx_split()
+                train_idx, val_idx, test_idx = (
+                    split["train"],
+                    split["valid"],
+                    split["test"],
+                )
+                n = self.g.num_nodes()
+                train_mask = utils.generate_mask_tensor(
+                    utils.idx2mask(train_idx, n)
+                )
+                val_mask = utils.generate_mask_tensor(
+                    utils.idx2mask(val_idx, n)
+                )
+                test_mask = utils.generate_mask_tensor(
+                    utils.idx2mask(test_idx, n)
+                )
+                self.g.ndata["train_mask"] = train_mask
+                self.g.ndata["val_mask"] = val_mask
+                self.g.ndata["test_mask"] = test_mask
+            else:
+                assert (
+                    "train_mask" in self.g.nodes[self.target_ntype].data
+                ), "train_mask is not provided, please specify split_ratio to generate the masks"
+                assert (
+                    "val_mask" in self.g.nodes[self.target_ntype].data
+                ), "val_mask is not provided, please specify split_ratio to generate the masks"
+                assert (
+                    "test_mask" in self.g.nodes[self.target_ntype].data
+                ), "test_mask is not provided, please specify split_ratio to generate the masks"
+        else:
+            if self.verbose:
+                print("Generating train/val/test masks...")
+            utils.add_nodepred_split(self, self.split_ratio, self.target_ntype)
+
+        self._set_split_index()
+
+        self.num_classes = getattr(self.dataset, "num_classes", None)
+        if self.num_classes is None:
+            self.num_classes = len(
+                F.unique(self.g.nodes[self.target_ntype].data["label"])
+            )
+
+    def has_cache(self):
+        return os.path.isfile(
+            os.path.join(self.save_path, "graph_{}.bin".format(self.hash))
+        )
+
+    def load(self):
+        with open(
+            os.path.join(self.save_path, "info_{}.json".format(self.hash)), "r"
+        ) as f:
+            info = json.load(f)
+            if (
+                info["split_ratio"] != self.split_ratio
+                or info["target_ntype"] != self.target_ntype
+            ):
+                raise ValueError(
+                    "Provided split ratio is different from the cached file. "
+                    "Re-process the dataset."
+                )
+            self.split_ratio = info["split_ratio"]
+            self.target_ntype = info["target_ntype"]
+            self.num_classes = info["num_classes"]
+        gs, _ = utils.load_graphs(
+            os.path.join(self.save_path, "graph_{}.bin".format(self.hash))
+        )
+        self.g = gs[0]
+        self._set_split_index()
+
+    def save(self):
+        utils.save_graphs(
+            os.path.join(self.save_path, "graph_{}.bin".format(self.hash)),
+            [self.g],
+        )
+        with open(
+            os.path.join(self.save_path, "info_{}.json".format(self.hash)), "w"
+        ) as f:
+            json.dump(
+                {
+                    "split_ratio": self.split_ratio,
+                    "target_ntype": self.target_ntype,
+                    "num_classes": self.num_classes,
+                },
+                f,
+            )
+
+    def __getitem__(self, idx):
+        return self.g
+
+    def __len__(self):
+        return 1
+
+    def _set_split_index(self):
+        """Add train_idx/val_idx/test_idx as dataset attributes according to corresponding mask."""
+        ndata = self.g.nodes[self.target_ntype].data
+        self.train_idx = F.nonzero_1d(ndata["train_mask"])
+        self.val_idx = F.nonzero_1d(ndata["val_mask"])
+        self.test_idx = F.nonzero_1d(ndata["test_mask"])
+
+
+def negative_sample(g, num_samples):
+    """Random sample negative edges from graph, excluding self-loops,
+    the result samples might be less than num_samples
+    """
+    num_nodes = g.num_nodes()
+    redundancy = _calc_redundancy(num_samples, g.num_edges(), num_nodes**2)
+    sample_size = int(num_samples * (1 + redundancy))
+    edges = np.random.randint(0, num_nodes, size=(2, sample_size))
+    edges = np.unique(edges, axis=1)
+    # remove self loop
+    mask_self_loop = edges[0] == edges[1]
+    # remove existing edges
+    has_edges = F.asnumpy(g.has_edges_between(edges[0], edges[1]))
+    mask = ~(np.logical_or(mask_self_loop, has_edges))
+    edges = edges[:, mask]
+    if edges.shape[1] >= num_samples:
+        edges = edges[:, :num_samples]
+    return edges
+
+
+class AsLinkPredDataset(DGLDataset):
+    """Repurpose a dataset for link prediction task.
+
+    The created dataset will include data needed for link prediction.
+    Currently it only supports homogeneous graphs.
+    It will keep only the first graph in the provided dataset and
+    generate train/val/test edges according to the given split ratio,
+    and the correspondent negative edges based on the neg_ratio. The generated
+    edges will be cached to disk for fast re-loading. If the provided split ratio
+    differs from the cached one, it will re-process the dataset properly.
+
+    Parameters
+    ----------
+    dataset : DGLDataset
+        The dataset to be converted.
+    split_ratio : (float, float, float), optional
+        Split ratios for training, validation and test sets. Must sum to one.
+    neg_ratio : int, optional
+        Indicate how much negative samples to be sampled
+        The number of the negative samples will be equal or less than neg_ratio * num_positive_edges.
+
+    Attributes
+    -------
+    feat_size: int
+        The size of the feature dimension in the graph
+    train_graph: DGLGraph
+        The DGLGraph for training
+    val_edges: Tuple[Tuple[Tensor, Tensor], Tuple[Tensor, Tensor]]
+        The validation set edges, encoded as
+        ((positive_edge_src, positive_edge_dst), (negative_edge_src, negative_edge_dst))
+    test_edges: Tuple[Tuple[Tensor, Tensor], Tuple[Tensor, Tensor]]
+        The test set edges, encoded as
+        ((positive_edge_src, positive_edge_dst), (negative_edge_src, negative_edge_dst))
+
+    Examples
+    --------
+    >>> ds = dgl.data.CoraGraphDataset()
+    >>> print(ds)
+    Dataset("cora_v2", num_graphs=1, save_path=...)
+    >>> new_ds = dgl.data.AsLinkPredDataset(ds, [0.8, 0.1, 0.1])
+    >>> print(new_ds)
+    Dataset("cora_v2-as-linkpred", num_graphs=1, save_path=/home/ubuntu/.dgl/cora_v2-as-linkpred)
+    >>> print(hasattr(new_ds, "test_edges"))
+    True
+    """
+
+    def __init__(self, dataset, split_ratio=None, neg_ratio=3, **kwargs):
+        self.g = dataset[0]
+        self.num_nodes = self.g.num_nodes()
+        self.dataset = dataset
+        self.split_ratio = split_ratio
+        self.neg_ratio = neg_ratio
+        super().__init__(
+            dataset.name + "-as-linkpred",
+            hash_key=(neg_ratio, split_ratio, dataset.name, "linkpred"),
+            **kwargs
+        )
+
+    def process(self):
+        if self.split_ratio is None:
+            # Handle logics for OGB link prediction dataset
+            assert hasattr(
+                self.dataset, "get_edge_split"
+            ), "dataset doesn't have get_edge_split method, please specify split_ratio and neg_ratio to generate the split"
+            # This is likely to be an ogb dataset
+            self.edge_split = self.dataset.get_edge_split()
+            self._train_graph = self.g
+            if "source_node" in self.edge_split["test"]:
+                # Probably ogbl-citation2
+                pos_e = (
+                    self.edge_split["valid"]["source_node"],
+                    self.edge_split["valid"]["target_node"],
+                )
+                neg_e_size = self.edge_split["valid"]["target_node_neg"].shape[
+                    -1
+                ]
+                neg_e_src = np.repeat(
+                    self.edge_split["valid"]["source_node"], neg_e_size
+                )
+                neg_e_dst = np.reshape(
+                    self.edge_split["valid"]["target_node_neg"], -1
+                )
+                self._val_edges = pos_e, (neg_e_src, neg_e_dst)
+                pos_e = (
+                    self.edge_split["test"]["source_node"],
+                    self.edge_split["test"]["target_node"],
+                )
+                neg_e_size = self.edge_split["test"]["target_node_neg"].shape[
+                    -1
+                ]
+                neg_e_src = np.repeat(
+                    self.edge_split["test"]["source_node"], neg_e_size
+                )
+                neg_e_dst = np.reshape(
+                    self.edge_split["test"]["target_node_neg"], -1
+                )
+                self._test_edges = pos_e, (neg_e_src, neg_e_dst)
+            elif "edge" in self.edge_split["test"]:
+                # Probably ogbl-collab
+                pos_e_tensor, neg_e_tensor = (
+                    self.edge_split["valid"]["edge"],
+                    self.edge_split["valid"]["edge_neg"],
+                )
+                pos_e = (pos_e_tensor[:, 0], pos_e_tensor[:, 1])
+                neg_e = (neg_e_tensor[:, 0], neg_e_tensor[:, 1])
+                self._val_edges = pos_e, neg_e
+
+                pos_e_tensor, neg_e_tensor = (
+                    self.edge_split["test"]["edge"],
+                    self.edge_split["test"]["edge_neg"],
+                )
+                pos_e = (pos_e_tensor[:, 0], pos_e_tensor[:, 1])
+                neg_e = (neg_e_tensor[:, 0], neg_e_tensor[:, 1])
+                self._test_edges = pos_e, neg_e
+            # delete edge split to save memory
+            self.edge_split = None
+        else:
+            assert self.split_ratio is not None, "Need to specify split_ratio"
+            assert self.neg_ratio is not None, "Need to specify neg_ratio"
+            ratio = self.split_ratio
+            graph = self.dataset[0]
+            n = graph.num_edges()
+            src, dst = graph.edges()
+            src, dst = F.asnumpy(src), F.asnumpy(dst)
+            n_train, n_val, n_test = (
+                int(n * ratio[0]),
+                int(n * ratio[1]),
+                int(n * ratio[2]),
+            )
+
+            idx = np.random.permutation(n)
+            train_pos_idx = idx[:n_train]
+            val_pos_idx = idx[n_train : n_train + n_val]
+            test_pos_idx = idx[n_train + n_val :]
+            neg_src, neg_dst = negative_sample(
+                graph, self.neg_ratio * (n_val + n_test)
+            )
+            neg_n_val, neg_n_test = (
+                self.neg_ratio * n_val,
+                self.neg_ratio * n_test,
+            )
+            neg_val_src, neg_val_dst = neg_src[:neg_n_val], neg_dst[:neg_n_val]
+            neg_test_src, neg_test_dst = (
+                neg_src[neg_n_val:],
+                neg_dst[neg_n_val:],
+            )
+            self._val_edges = (
+                F.tensor(src[val_pos_idx]),
+                F.tensor(dst[val_pos_idx]),
+            ), (F.tensor(neg_val_src), F.tensor(neg_val_dst))
+            self._test_edges = (
+                F.tensor(src[test_pos_idx]),
+                F.tensor(dst[test_pos_idx]),
+            ), (F.tensor(neg_test_src), F.tensor(neg_test_dst))
+            self._train_graph = create_dgl_graph(
+                (src[train_pos_idx], dst[train_pos_idx]),
+                num_nodes=self.num_nodes,
+            )
+            self._train_graph.ndata["feat"] = graph.ndata["feat"]
+
+    def has_cache(self):
+        return os.path.isfile(
+            os.path.join(self.save_path, "graph_{}.bin".format(self.hash))
+        )
+
+    def load(self):
+        gs, tensor_dict = utils.load_graphs(
+            os.path.join(self.save_path, "graph_{}.bin".format(self.hash))
+        )
+        self.g = gs[0]
+        self._train_graph = self.g
+        self._val_edges = (
+            tensor_dict["val_pos_src"],
+            tensor_dict["val_pos_dst"],
+        ), (tensor_dict["val_neg_src"], tensor_dict["val_neg_dst"])
+        self._test_edges = (
+            tensor_dict["test_pos_src"],
+            tensor_dict["test_pos_dst"],
+        ), (tensor_dict["test_neg_src"], tensor_dict["test_neg_dst"])
+
+        with open(
+            os.path.join(self.save_path, "info_{}.json".format(self.hash)), "r"
+        ) as f:
+            info = json.load(f)
+            self.split_ratio = info["split_ratio"]
+            self.neg_ratio = info["neg_ratio"]
+
+    def save(self):
+        tensor_dict = {
+            "val_pos_src": self._val_edges[0][0],
+            "val_pos_dst": self._val_edges[0][1],
+            "val_neg_src": self._val_edges[1][0],
+            "val_neg_dst": self._val_edges[1][1],
+            "test_pos_src": self._test_edges[0][0],
+            "test_pos_dst": self._test_edges[0][1],
+            "test_neg_src": self._test_edges[1][0],
+            "test_neg_dst": self._test_edges[1][1],
+        }
+        utils.save_graphs(
+            os.path.join(self.save_path, "graph_{}.bin".format(self.hash)),
+            [self._train_graph],
+            tensor_dict,
+        )
+        with open(
+            os.path.join(self.save_path, "info_{}.json".format(self.hash)), "w"
+        ) as f:
+            json.dump(
+                {"split_ratio": self.split_ratio, "neg_ratio": self.neg_ratio},
+                f,
+            )
+
+    @property
+    def feat_size(self):
+        return self._train_graph.ndata["feat"].shape[-1]
+
+    @property
+    def train_graph(self):
+        return self._train_graph
+
+    @property
+    def val_edges(self):
+        return self._val_edges
+
+    @property
+    def test_edges(self):
+        return self._test_edges
+
+    def __getitem__(self, idx):
+        return self.g
+
+    def __len__(self):
+        return 1
+
+
+class AsGraphPredDataset(DGLDataset):
+    """Repurpose a dataset for standard graph property prediction task.
+
+    The created dataset will include data needed for graph property prediction.
+    Currently it only supports homogeneous graphs.
+
+    The class converts a given dataset into a new dataset object such that:
+
+      - It stores ``len(dataset)`` graphs.
+      - The i-th graph and its label is accessible from ``dataset[i]``.
+
+    The class will generate a train/val/test split if :attr:`split_ratio` is provided.
+    The generated split will be cached to disk for fast re-loading. If the provided split
+    ratio differs from the cached one, it will re-process the dataset properly.
+
+    Parameters
+    ----------
+    dataset : DGLDataset
+        The dataset to be converted.
+    split_ratio : (float, float, float), optional
+        Split ratios for training, validation and test sets. They must sum to one.
+
+    Attributes
+    ----------
+    num_tasks : int
+        Number of tasks to predict.
+    num_classes : int
+        Number of classes to predict per task, None for regression datasets.
+    train_idx : Tensor
+        An 1-D integer tensor of training node IDs.
+    val_idx : Tensor
+        An 1-D integer tensor of validation node IDs.
+    test_idx : Tensor
+        An 1-D integer tensor of test node IDs.
+    node_feat_size : int
+        Input node feature size, None if not applicable.
+    edge_feat_size : int
+        Input edge feature size, None if not applicable.
+
+    Examples
+    --------
+
+    >>> from dgl.data import AsGraphPredDataset
+    >>> from ogb.graphproppred import DglGraphPropPredDataset
+    >>> dataset = DglGraphPropPredDataset(name='ogbg-molhiv')
+    >>> new_dataset = AsGraphPredDataset(dataset)
+    >>> print(new_dataset)
+    Dataset("ogbg-molhiv-as-graphpred", num_graphs=41127, save_path=...)
+    >>> print(len(new_dataset))
+    41127
+    >>> print(new_dataset[0])
+    (Graph(num_nodes=19, num_edges=40,
+           ndata_schemes={'feat': Scheme(shape=(9,), dtype=torch.int64)}
+           edata_schemes={'feat': Scheme(shape=(3,), dtype=torch.int64)}), tensor([0]))
+    """
+
+    def __init__(self, dataset, split_ratio=None, **kwargs):
+        self.dataset = dataset
+        self.split_ratio = split_ratio
+        super().__init__(
+            dataset.name + "-as-graphpred",
+            hash_key=(split_ratio, dataset.name, "graphpred"),
+            **kwargs
+        )
+
+    def process(self):
+        is_ogb = hasattr(self.dataset, "get_idx_split")
+        if self.split_ratio is None:
+            if is_ogb:
+                split = self.dataset.get_idx_split()
+                self.train_idx = split["train"]
+                self.val_idx = split["valid"]
+                self.test_idx = split["test"]
+            else:
+                # Handle FakeNewsDataset
+                try:
+                    self.train_idx = F.nonzero_1d(self.dataset.train_mask)
+                    self.val_idx = F.nonzero_1d(self.dataset.val_mask)
+                    self.test_idx = F.nonzero_1d(self.dataset.test_mask)
+                except:
+                    raise DGLError(
+                        "The input dataset does not have default train/val/test\
+                        split. Please specify split_ratio to generate the split."
+                    )
+        else:
+            if self.verbose:
+                print("Generating train/val/test split...")
+            train_ratio, val_ratio, _ = self.split_ratio
+            num_graphs = len(self.dataset)
+            num_train = int(num_graphs * train_ratio)
+            num_val = int(num_graphs * val_ratio)
+
+            idx = np.random.permutation(num_graphs)
+            self.train_idx = F.tensor(idx[:num_train])
+            self.val_idx = F.tensor(idx[num_train : num_train + num_val])
+            self.test_idx = F.tensor(idx[num_train + num_val :])
+
+        if hasattr(self.dataset, "num_classes"):
+            # GINDataset, MiniGCDataset, FakeNewsDataset, TUDataset,
+            # LegacyTUDataset, BA2MotifDataset
+            self.num_classes = self.dataset.num_classes
+        else:
+            # None for multi-label classification and regression
+            self.num_classes = None
+
+        if hasattr(self.dataset, "num_tasks"):
+            # OGB datasets
+            self.num_tasks = self.dataset.num_tasks
+        else:
+            self.num_tasks = 1
+
+    def has_cache(self):
+        return os.path.isfile(
+            os.path.join(self.save_path, "info_{}.json".format(self.hash))
+        )
+
+    def load(self):
+        with open(
+            os.path.join(self.save_path, "info_{}.json".format(self.hash)), "r"
+        ) as f:
+            info = json.load(f)
+            if info["split_ratio"] != self.split_ratio:
+                raise ValueError(
+                    "Provided split ratio is different from the cached file. "
+                    "Re-process the dataset."
+                )
+            self.split_ratio = info["split_ratio"]
+            self.num_tasks = info["num_tasks"]
+            self.num_classes = info["num_classes"]
+
+        split = np.load(
+            os.path.join(self.save_path, "split_{}.npz".format(self.hash))
+        )
+        self.train_idx = F.zerocopy_from_numpy(split["train_idx"])
+        self.val_idx = F.zerocopy_from_numpy(split["val_idx"])
+        self.test_idx = F.zerocopy_from_numpy(split["test_idx"])
+
+    def save(self):
+        if not os.path.exists(self.save_path):
+            os.makedirs(self.save_path)
+        with open(
+            os.path.join(self.save_path, "info_{}.json".format(self.hash)), "w"
+        ) as f:
+            json.dump(
+                {
+                    "split_ratio": self.split_ratio,
+                    "num_tasks": self.num_tasks,
+                    "num_classes": self.num_classes,
+                },
+                f,
+            )
+        np.savez(
+            os.path.join(self.save_path, "split_{}.npz".format(self.hash)),
+            train_idx=F.zerocopy_to_numpy(self.train_idx),
+            val_idx=F.zerocopy_to_numpy(self.val_idx),
+            test_idx=F.zerocopy_to_numpy(self.test_idx),
+        )
+
+    def __getitem__(self, idx):
+        return self.dataset[idx]
+
+    def __len__(self):
+        return len(self.dataset)
+
+    @property
+    def node_feat_size(self):
+        g = self[0][0]
+        return g.ndata["feat"].shape[-1] if "feat" in g.ndata else None
+
+    @property
+    def edge_feat_size(self):
+        g = self[0][0]
+        return g.edata["feat"].shape[-1] if "feat" in g.edata else None

python/dgl/data/bitcoinotc.py

@@ -0,0 +1,191 @@
+""" BitcoinOTC dataset for fraud detection """
+import datetime
+import gzip
+import os
+import shutil
+
+import numpy as np
+
+from .. import backend as F
+from ..convert import graph as dgl_graph
+from .dgl_dataset import DGLBuiltinDataset
+from .utils import check_sha1, download, load_graphs, makedirs, save_graphs
+
+
+class BitcoinOTCDataset(DGLBuiltinDataset):
+    r"""BitcoinOTC dataset for fraud detection
+
+    This is who-trusts-whom network of people who trade using Bitcoin on
+    a platform called Bitcoin OTC. Since Bitcoin users are anonymous,
+    there is a need to maintain a record of users' reputation to prevent
+    transactions with fraudulent and risky users.
+
+    Offical website: `<https://snap.stanford.edu/data/soc-sign-bitcoin-otc.html>`_
+
+    Bitcoin OTC dataset statistics:
+
+    - Nodes: 5,881
+    - Edges: 35,592
+    - Range of edge weight: -10 to +10
+    - Percentage of positive edges: 89%
+
+    Parameters
+    ----------
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset.
+        Default: False
+    verbose: bool
+        Whether to print out progress information.
+        Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    graphs : list
+        A list of DGLGraph objects
+    is_temporal : bool
+        Indicate whether the graphs are temporal graphs
+
+    Raises
+    ------
+    UserWarning
+        If the raw data is changed in the remote server by the author.
+
+    Examples
+    --------
+    >>> dataset = BitcoinOTCDataset()
+    >>> len(dataset)
+    136
+    >>> for g in dataset:
+    ....    # get edge feature
+    ....    edge_weights = g.edata['h']
+    ....    # your code here
+    >>>
+    """
+
+    _url = "https://snap.stanford.edu/data/soc-sign-bitcoinotc.csv.gz"
+    _sha1_str = "c14281f9e252de0bd0b5f1c6e2bae03123938641"
+
+    def __init__(
+        self, raw_dir=None, force_reload=False, verbose=False, transform=None
+    ):
+        super(BitcoinOTCDataset, self).__init__(
+            name="bitcoinotc",
+            url=self._url,
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def download(self):
+        gz_file_path = os.path.join(self.raw_dir, self.name + ".csv.gz")
+        download(self.url, path=gz_file_path)
+        if not check_sha1(gz_file_path, self._sha1_str):
+            raise UserWarning(
+                "File {} is downloaded but the content hash does not match."
+                "The repo may be outdated or download may be incomplete. "
+                "Otherwise you can create an issue for it.".format(
+                    self.name + ".csv.gz"
+                )
+            )
+        self._extract_gz(gz_file_path, self.raw_path)
+
+    def process(self):
+        filename = os.path.join(self.save_path, self.name + ".csv")
+        data = np.loadtxt(filename, delimiter=",").astype(np.int64)
+        data[:, 0:2] = data[:, 0:2] - data[:, 0:2].min()
+        delta = datetime.timedelta(days=14).total_seconds()
+        # The source code is not released, but the paper indicates there're
+        # totally 137 samples. The cutoff below has exactly 137 samples.
+        time_index = np.around((data[:, 3] - data[:, 3].min()) / delta).astype(
+            np.int64
+        )
+
+        self._graphs = []
+        for i in range(time_index.max()):
+            row_mask = time_index <= i
+            edges = data[row_mask][:, 0:2]
+            rate = data[row_mask][:, 2]
+            g = dgl_graph((edges[:, 0], edges[:, 1]))
+            g.edata["h"] = F.tensor(
+                rate.reshape(-1, 1), dtype=F.data_type_dict["int64"]
+            )
+            self._graphs.append(g)
+
+    @property
+    def graph_path(self):
+        return os.path.join(self.save_path, "dgl_graph.bin")
+
+    def has_cache(self):
+        return os.path.exists(self.graph_path)
+
+    def save(self):
+        save_graphs(self.graph_path, self.graphs)
+
+    def load(self):
+        self._graphs = load_graphs(self.graph_path)[0]
+
+    @property
+    def graphs(self):
+        return self._graphs
+
+    def __len__(self):
+        r"""Number of graphs in the dataset.
+
+        Return
+        -------
+        int
+        """
+        return len(self.graphs)
+
+    def __getitem__(self, item):
+        r"""Get graph by index
+
+        Parameters
+        ----------
+        item : int
+            Item index
+
+        Returns
+        -------
+        :class:`dgl.DGLGraph`
+
+            The graph contains:
+
+            - ``edata['h']`` : edge weights
+        """
+        if self._transform is None:
+            return self.graphs[item]
+        else:
+            return self._transform(self.graphs[item])
+
+    @property
+    def is_temporal(self):
+        r"""Are the graphs temporal graphs
+
+        Returns
+        -------
+        bool
+        """
+        return True
+
+    def _extract_gz(self, file, target_dir, overwrite=False):
+        if os.path.exists(target_dir) and not overwrite:
+            return
+        print("Extracting file to {}".format(target_dir))
+        fname = os.path.basename(file)
+        makedirs(target_dir)
+        out_file_path = os.path.join(target_dir, fname[:-3])
+        with gzip.open(file, "rb") as f_in:
+            with open(out_file_path, "wb") as f_out:
+                shutil.copyfileobj(f_in, f_out)
+
+
+BitcoinOTC = BitcoinOTCDataset

python/dgl/data/citation_graph.py

@@ -0,0 +1,953 @@
+"""Cora, citeseer, pubmed dataset.
+
+(lingfan): following dataset loading and preprocessing code from tkipf/gcn
+https://github.com/tkipf/gcn/blob/master/gcn/utils.py
+"""
+
+from __future__ import absolute_import
+
+import os, sys
+import pickle as pkl
+import warnings
+
+import networkx as nx
+
+import numpy as np
+import scipy.sparse as sp
+
+from .. import backend as F, convert
+from ..batch import batch as batch_graphs
+from ..convert import from_networkx, graph as dgl_graph, to_networkx
+from ..transforms import reorder_graph
+from .dgl_dataset import DGLBuiltinDataset
+
+from .utils import (
+    _get_dgl_url,
+    deprecate_function,
+    deprecate_property,
+    generate_mask_tensor,
+    load_graphs,
+    load_info,
+    makedirs,
+    save_graphs,
+    save_info,
+)
+
+backend = os.environ.get("DGLBACKEND", "pytorch")
+
+
+def _pickle_load(pkl_file):
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", category=DeprecationWarning)
+        if sys.version_info > (3, 0):
+            return pkl.load(pkl_file, encoding="latin1")
+        else:
+            return pkl.load(pkl_file)
+
+
+class CitationGraphDataset(DGLBuiltinDataset):
+    r"""The citation graph dataset, including cora, citeseer and pubmeb.
+    Nodes mean authors and edges mean citation relationships.
+
+    Parameters
+    -----------
+    name: str
+      name can be 'cora', 'citeseer' or 'pubmed'.
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    reverse_edge : bool
+        Whether to add reverse edges in graph. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+    reorder : bool
+        Whether to reorder the graph using :func:`~dgl.reorder_graph`. Default: False.
+    """
+
+    _urls = {
+        "cora_v2": "dataset/cora_v2.zip",
+        "citeseer": "dataset/citeseer.zip",
+        "pubmed": "dataset/pubmed.zip",
+    }
+
+    def __init__(
+        self,
+        name,
+        raw_dir=None,
+        force_reload=False,
+        verbose=True,
+        reverse_edge=True,
+        transform=None,
+        reorder=False,
+    ):
+        assert name.lower() in ["cora", "citeseer", "pubmed"]
+
+        # Previously we use the pre-processing in pygcn (https://github.com/tkipf/pygcn)
+        # for Cora, which is slightly different from the one used in the GCN paper
+        if name.lower() == "cora":
+            name = "cora_v2"
+
+        url = _get_dgl_url(self._urls[name])
+        self._reverse_edge = reverse_edge
+        self._reorder = reorder
+
+        super(CitationGraphDataset, self).__init__(
+            name,
+            url=url,
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        """Loads input data from data directory and reorder graph for better locality
+
+        ind.name.x => the feature vectors of the training instances as scipy.sparse.csr.csr_matrix object;
+        ind.name.tx => the feature vectors of the test instances as scipy.sparse.csr.csr_matrix object;
+        ind.name.allx => the feature vectors of both labeled and unlabeled training instances
+            (a superset of ind.name.x) as scipy.sparse.csr.csr_matrix object;
+        ind.name.y => the one-hot labels of the labeled training instances as numpy.ndarray object;
+        ind.name.ty => the one-hot labels of the test instances as numpy.ndarray object;
+        ind.name.ally => the labels for instances in ind.name.allx as numpy.ndarray object;
+        ind.name.graph => a dict in the format {index: [index_of_neighbor_nodes]} as collections.defaultdict
+            object;
+        ind.name.test.index => the indices of test instances in graph, for the inductive setting as list object.
+        """
+        root = self.raw_path
+        objnames = ["x", "y", "tx", "ty", "allx", "ally", "graph"]
+        objects = []
+        for i in range(len(objnames)):
+            with open(
+                "{}/ind.{}.{}".format(root, self.name, objnames[i]), "rb"
+            ) as f:
+                objects.append(_pickle_load(f))
+
+        x, y, tx, ty, allx, ally, graph = tuple(objects)
+        test_idx_reorder = _parse_index_file(
+            "{}/ind.{}.test.index".format(root, self.name)
+        )
+        test_idx_range = np.sort(test_idx_reorder)
+
+        if self.name == "citeseer":
+            # Fix citeseer dataset (there are some isolated nodes in the graph)
+            # Find isolated nodes, add them as zero-vecs into the right position
+            test_idx_range_full = range(
+                min(test_idx_reorder), max(test_idx_reorder) + 1
+            )
+            tx_extended = sp.lil_matrix((len(test_idx_range_full), x.shape[1]))
+            tx_extended[test_idx_range - min(test_idx_range), :] = tx
+            tx = tx_extended
+            ty_extended = np.zeros((len(test_idx_range_full), y.shape[1]))
+            ty_extended[test_idx_range - min(test_idx_range), :] = ty
+            ty = ty_extended
+
+        features = sp.vstack((allx, tx)).tolil()
+        features[test_idx_reorder, :] = features[test_idx_range, :]
+
+        if self.reverse_edge:
+            graph = nx.DiGraph(nx.from_dict_of_lists(graph))
+            g = from_networkx(graph)
+        else:
+            graph = nx.Graph(nx.from_dict_of_lists(graph))
+            edges = list(graph.edges())
+            u, v = map(list, zip(*edges))
+            g = dgl_graph((u, v))
+
+        onehot_labels = np.vstack((ally, ty))
+        onehot_labels[test_idx_reorder, :] = onehot_labels[test_idx_range, :]
+        labels = np.argmax(onehot_labels, 1)
+
+        idx_test = test_idx_range.tolist()
+        idx_train = range(len(y))
+        idx_val = range(len(y), len(y) + 500)
+
+        train_mask = generate_mask_tensor(
+            _sample_mask(idx_train, labels.shape[0])
+        )
+        val_mask = generate_mask_tensor(_sample_mask(idx_val, labels.shape[0]))
+        test_mask = generate_mask_tensor(
+            _sample_mask(idx_test, labels.shape[0])
+        )
+
+        g.ndata["train_mask"] = train_mask
+        g.ndata["val_mask"] = val_mask
+        g.ndata["test_mask"] = test_mask
+        g.ndata["label"] = F.tensor(labels)
+        g.ndata["feat"] = F.tensor(
+            _preprocess_features(features), dtype=F.data_type_dict["float32"]
+        )
+        self._num_classes = onehot_labels.shape[1]
+        self._labels = labels
+        if self._reorder:
+            self._g = reorder_graph(
+                g,
+                node_permute_algo="rcmk",
+                edge_permute_algo="dst",
+                store_ids=False,
+            )
+        else:
+            self._g = g
+
+        if self.verbose:
+            print("Finished data loading and preprocessing.")
+            print("  NumNodes: {}".format(self._g.num_nodes()))
+            print("  NumEdges: {}".format(self._g.num_edges()))
+            print("  NumFeats: {}".format(self._g.ndata["feat"].shape[1]))
+            print("  NumClasses: {}".format(self.num_classes))
+            print(
+                "  NumTrainingSamples: {}".format(
+                    F.nonzero_1d(self._g.ndata["train_mask"]).shape[0]
+                )
+            )
+            print(
+                "  NumValidationSamples: {}".format(
+                    F.nonzero_1d(self._g.ndata["val_mask"]).shape[0]
+                )
+            )
+            print(
+                "  NumTestSamples: {}".format(
+                    F.nonzero_1d(self._g.ndata["test_mask"]).shape[0]
+                )
+            )
+
+    @property
+    def graph_path(self):
+        return os.path.join(self.save_path, self.save_name + ".bin")
+
+    @property
+    def info_path(self):
+        return os.path.join(self.save_path, self.save_name + ".pkl")
+
+    def has_cache(self):
+        if os.path.exists(self.graph_path) and os.path.exists(self.info_path):
+            return True
+
+        return False
+
+    def save(self):
+        """save the graph list and the labels"""
+        save_graphs(str(self.graph_path), self._g)
+        save_info(str(self.info_path), {"num_classes": self.num_classes})
+
+    def load(self):
+        graphs, _ = load_graphs(str(self.graph_path))
+
+        info = load_info(str(self.info_path))
+        graph = graphs[0]
+        self._g = graph
+        # for compatability
+        graph = graph.clone()
+        graph.ndata.pop("train_mask")
+        graph.ndata.pop("val_mask")
+        graph.ndata.pop("test_mask")
+        graph.ndata.pop("feat")
+        graph.ndata.pop("label")
+        graph = to_networkx(graph)
+
+        self._num_classes = info["num_classes"]
+        self._g.ndata["train_mask"] = generate_mask_tensor(
+            F.asnumpy(self._g.ndata["train_mask"])
+        )
+        self._g.ndata["val_mask"] = generate_mask_tensor(
+            F.asnumpy(self._g.ndata["val_mask"])
+        )
+        self._g.ndata["test_mask"] = generate_mask_tensor(
+            F.asnumpy(self._g.ndata["test_mask"])
+        )
+        # hack for mxnet compatability
+
+        if self.verbose:
+            print("  NumNodes: {}".format(self._g.num_nodes()))
+            print("  NumEdges: {}".format(self._g.num_edges()))
+            print("  NumFeats: {}".format(self._g.ndata["feat"].shape[1]))
+            print("  NumClasses: {}".format(self.num_classes))
+            print(
+                "  NumTrainingSamples: {}".format(
+                    F.nonzero_1d(self._g.ndata["train_mask"]).shape[0]
+                )
+            )
+            print(
+                "  NumValidationSamples: {}".format(
+                    F.nonzero_1d(self._g.ndata["val_mask"]).shape[0]
+                )
+            )
+            print(
+                "  NumTestSamples: {}".format(
+                    F.nonzero_1d(self._g.ndata["test_mask"]).shape[0]
+                )
+            )
+
+    def __getitem__(self, idx):
+        assert idx == 0, "This dataset has only one graph"
+        if self._transform is None:
+            return self._g
+        else:
+            return self._transform(self._g)
+
+    def __len__(self):
+        return 1
+
+    @property
+    def save_name(self):
+        return self.name + "_dgl_graph"
+
+    @property
+    def num_labels(self):
+        deprecate_property("dataset.num_labels", "dataset.num_classes")
+        return self.num_classes
+
+    @property
+    def num_classes(self):
+        return self._num_classes
+
+    """ Citation graph is used in many examples
+        We preserve these properties for compatability.
+    """
+
+    @property
+    def reverse_edge(self):
+        return self._reverse_edge
+
+
+def _preprocess_features(features):
+    """Row-normalize feature matrix and convert to tuple representation"""
+    features = _normalize(features)
+    return np.asarray(features.todense())
+
+
+def _parse_index_file(filename):
+    """Parse index file."""
+    index = []
+    for line in open(filename):
+        index.append(int(line.strip()))
+    return index
+
+
+def _sample_mask(idx, l):
+    """Create mask."""
+    mask = np.zeros(l)
+    mask[idx] = 1
+    return mask
+
+
+class CoraGraphDataset(CitationGraphDataset):
+    r"""Cora citation network dataset.
+
+    Nodes mean paper and edges mean citation
+    relationships. Each node has a predefined
+    feature with 1433 dimensions. The dataset is
+    designed for the node classification task.
+    The task is to predict the category of
+    certain paper.
+
+    Statistics:
+
+    - Nodes: 2708
+    - Edges: 10556
+    - Number of Classes: 7
+    - Label split:
+
+        - Train: 140
+        - Valid: 500
+        - Test: 1000
+
+    Parameters
+    ----------
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    reverse_edge : bool
+        Whether to add reverse edges in graph. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+    reorder : bool
+        Whether to reorder the graph using :func:`~dgl.reorder_graph`. Default: False.
+
+    Attributes
+    ----------
+    num_classes: int
+        Number of label classes
+
+    Notes
+    -----
+    The node feature is row-normalized.
+
+    Examples
+    --------
+    >>> dataset = CoraGraphDataset()
+    >>> g = dataset[0]
+    >>> num_class = dataset.num_classes
+    >>>
+    >>> # get node feature
+    >>> feat = g.ndata['feat']
+    >>>
+    >>> # get data split
+    >>> train_mask = g.ndata['train_mask']
+    >>> val_mask = g.ndata['val_mask']
+    >>> test_mask = g.ndata['test_mask']
+    >>>
+    >>> # get labels
+    >>> label = g.ndata['label']
+
+    """
+
+    def __init__(
+        self,
+        raw_dir=None,
+        force_reload=False,
+        verbose=True,
+        reverse_edge=True,
+        transform=None,
+        reorder=False,
+    ):
+        name = "cora"
+
+        super(CoraGraphDataset, self).__init__(
+            name,
+            raw_dir,
+            force_reload,
+            verbose,
+            reverse_edge,
+            transform,
+            reorder,
+        )
+
+    def __getitem__(self, idx):
+        r"""Gets the graph object
+
+        Parameters
+        -----------
+        idx: int
+            Item index, CoraGraphDataset has only one graph object
+
+        Return
+        ------
+        :class:`dgl.DGLGraph`
+
+            graph structure, node features and labels.
+
+            - ``ndata['train_mask']``: mask for training node set
+            - ``ndata['val_mask']``: mask for validation node set
+            - ``ndata['test_mask']``: mask for test node set
+            - ``ndata['feat']``: node feature
+            - ``ndata['label']``: ground truth labels
+        """
+        return super(CoraGraphDataset, self).__getitem__(idx)
+
+    def __len__(self):
+        r"""The number of graphs in the dataset."""
+        return super(CoraGraphDataset, self).__len__()
+
+
+class CiteseerGraphDataset(CitationGraphDataset):
+    r"""Citeseer citation network dataset.
+
+    Nodes mean scientific publications and edges
+    mean citation relationships. Each node has a
+    predefined feature with 3703 dimensions. The
+    dataset is designed for the node classification
+    task. The task is to predict the category of
+    certain publication.
+
+    Statistics:
+
+    - Nodes: 3327
+    - Edges: 9228
+    - Number of Classes: 6
+    - Label Split:
+
+        - Train: 120
+        - Valid: 500
+        - Test: 1000
+
+    Parameters
+    -----------
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    reverse_edge : bool
+        Whether to add reverse edges in graph. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+    reorder : bool
+        Whether to reorder the graph using :func:`~dgl.reorder_graph`. Default: False.
+
+    Attributes
+    ----------
+    num_classes: int
+        Number of label classes
+
+    Notes
+    -----
+    The node feature is row-normalized.
+
+    In citeseer dataset, there are some isolated nodes in the graph.
+    These isolated nodes are added as zero-vecs into the right position.
+
+    Examples
+    --------
+    >>> dataset = CiteseerGraphDataset()
+    >>> g = dataset[0]
+    >>> num_class = dataset.num_classes
+    >>>
+    >>> # get node feature
+    >>> feat = g.ndata['feat']
+    >>>
+    >>> # get data split
+    >>> train_mask = g.ndata['train_mask']
+    >>> val_mask = g.ndata['val_mask']
+    >>> test_mask = g.ndata['test_mask']
+    >>>
+    >>> # get labels
+    >>> label = g.ndata['label']
+
+    """
+
+    def __init__(
+        self,
+        raw_dir=None,
+        force_reload=False,
+        verbose=True,
+        reverse_edge=True,
+        transform=None,
+        reorder=False,
+    ):
+        name = "citeseer"
+
+        super(CiteseerGraphDataset, self).__init__(
+            name,
+            raw_dir,
+            force_reload,
+            verbose,
+            reverse_edge,
+            transform,
+            reorder,
+        )
+
+    def __getitem__(self, idx):
+        r"""Gets the graph object
+
+        Parameters
+        -----------
+        idx: int
+            Item index, CiteseerGraphDataset has only one graph object
+
+        Return
+        ------
+        :class:`dgl.DGLGraph`
+
+            graph structure, node features and labels.
+
+            - ``ndata['train_mask']``: mask for training node set
+            - ``ndata['val_mask']``: mask for validation node set
+            - ``ndata['test_mask']``: mask for test node set
+            - ``ndata['feat']``: node feature
+            - ``ndata['label']``: ground truth labels
+        """
+        return super(CiteseerGraphDataset, self).__getitem__(idx)
+
+    def __len__(self):
+        r"""The number of graphs in the dataset."""
+        return super(CiteseerGraphDataset, self).__len__()
+
+
+class PubmedGraphDataset(CitationGraphDataset):
+    r"""Pubmed citation network dataset.
+
+    Nodes mean scientific publications and edges
+    mean citation relationships. Each node has a
+    predefined feature with 500 dimensions. The
+    dataset is designed for the node classification
+    task. The task is to predict the category of
+    certain publication.
+
+    Statistics:
+
+    - Nodes: 19717
+    - Edges: 88651
+    - Number of Classes: 3
+    - Label Split:
+
+        - Train: 60
+        - Valid: 500
+        - Test: 1000
+
+    Parameters
+    -----------
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    reverse_edge : bool
+        Whether to add reverse edges in graph. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+    reorder : bool
+        Whether to reorder the graph using :func:`~dgl.reorder_graph`. Default: False.
+
+    Attributes
+    ----------
+    num_classes: int
+        Number of label classes
+
+    Notes
+    -----
+    The node feature is row-normalized.
+
+    Examples
+    --------
+    >>> dataset = PubmedGraphDataset()
+    >>> g = dataset[0]
+    >>> num_class = dataset.num_of_class
+    >>>
+    >>> # get node feature
+    >>> feat = g.ndata['feat']
+    >>>
+    >>> # get data split
+    >>> train_mask = g.ndata['train_mask']
+    >>> val_mask = g.ndata['val_mask']
+    >>> test_mask = g.ndata['test_mask']
+    >>>
+    >>> # get labels
+    >>> label = g.ndata['label']
+
+    """
+
+    def __init__(
+        self,
+        raw_dir=None,
+        force_reload=False,
+        verbose=True,
+        reverse_edge=True,
+        transform=None,
+        reorder=False,
+    ):
+        name = "pubmed"
+
+        super(PubmedGraphDataset, self).__init__(
+            name,
+            raw_dir,
+            force_reload,
+            verbose,
+            reverse_edge,
+            transform,
+            reorder,
+        )
+
+    def __getitem__(self, idx):
+        r"""Gets the graph object
+
+        Parameters
+        -----------
+        idx: int
+            Item index, PubmedGraphDataset has only one graph object
+
+        Return
+        ------
+        :class:`dgl.DGLGraph`
+
+            graph structure, node features and labels.
+
+            - ``ndata['train_mask']``: mask for training node set
+            - ``ndata['val_mask']``: mask for validation node set
+            - ``ndata['test_mask']``: mask for test node set
+            - ``ndata['feat']``: node feature
+            - ``ndata['label']``: ground truth labels
+        """
+        return super(PubmedGraphDataset, self).__getitem__(idx)
+
+    def __len__(self):
+        r"""The number of graphs in the dataset."""
+        return super(PubmedGraphDataset, self).__len__()
+
+
+def load_cora(
+    raw_dir=None,
+    force_reload=False,
+    verbose=True,
+    reverse_edge=True,
+    transform=None,
+):
+    """Get CoraGraphDataset
+
+    Parameters
+    -----------
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    reverse_edge : bool
+        Whether to add reverse edges in graph. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Return
+    -------
+    CoraGraphDataset
+    """
+    data = CoraGraphDataset(
+        raw_dir, force_reload, verbose, reverse_edge, transform
+    )
+    return data
+
+
+def load_citeseer(
+    raw_dir=None,
+    force_reload=False,
+    verbose=True,
+    reverse_edge=True,
+    transform=None,
+):
+    """Get CiteseerGraphDataset
+
+    Parameters
+    -----------
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    reverse_edge : bool
+        Whether to add reverse edges in graph. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Return
+    -------
+    CiteseerGraphDataset
+    """
+    data = CiteseerGraphDataset(
+        raw_dir, force_reload, verbose, reverse_edge, transform
+    )
+    return data
+
+
+def load_pubmed(
+    raw_dir=None,
+    force_reload=False,
+    verbose=True,
+    reverse_edge=True,
+    transform=None,
+):
+    """Get PubmedGraphDataset
+
+    Parameters
+    -----------
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    reverse_edge : bool
+        Whether to add reverse edges in graph. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Return
+    -------
+    PubmedGraphDataset
+    """
+    data = PubmedGraphDataset(
+        raw_dir, force_reload, verbose, reverse_edge, transform
+    )
+    return data
+
+
+class CoraBinary(DGLBuiltinDataset):
+    """A mini-dataset for binary classification task using Cora.
+
+    After loaded, it has following members:
+
+    graphs : list of :class:`~dgl.DGLGraph`
+    pmpds : list of :class:`scipy.sparse.coo_matrix`
+    labels : list of :class:`numpy.ndarray`
+
+    Parameters
+    -----------
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose: bool
+        Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+    """
+
+    def __init__(
+        self, raw_dir=None, force_reload=False, verbose=True, transform=None
+    ):
+        name = "cora_binary"
+        url = _get_dgl_url("dataset/cora_binary.zip")
+        super(CoraBinary, self).__init__(
+            name,
+            url=url,
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        root = self.raw_path
+        # load graphs
+        self.graphs = []
+        with open("{}/graphs.txt".format(root), "r") as f:
+            elist = []
+            for line in f.readlines():
+                if line.startswith("graph"):
+                    if len(elist) != 0:
+                        self.graphs.append(dgl_graph(tuple(zip(*elist))))
+                    elist = []
+                else:
+                    u, v = line.strip().split(" ")
+                    elist.append((int(u), int(v)))
+            if len(elist) != 0:
+                self.graphs.append(dgl_graph(tuple(zip(*elist))))
+        with open("{}/pmpds.pkl".format(root), "rb") as f:
+            self.pmpds = _pickle_load(f)
+        self.labels = []
+        with open("{}/labels.txt".format(root), "r") as f:
+            cur = []
+            for line in f.readlines():
+                if line.startswith("graph"):
+                    if len(cur) != 0:
+                        self.labels.append(np.asarray(cur))
+                    cur = []
+                else:
+                    cur.append(int(line.strip()))
+            if len(cur) != 0:
+                self.labels.append(np.asarray(cur))
+        # sanity check
+        assert len(self.graphs) == len(self.pmpds)
+        assert len(self.graphs) == len(self.labels)
+
+    @property
+    def graph_path(self):
+        return os.path.join(self.save_path, self.save_name + ".bin")
+
+    def has_cache(self):
+        if os.path.exists(self.graph_path):
+            return True
+
+        return False
+
+    def save(self):
+        """save the graph list and the labels"""
+        labels = {}
+        for i, label in enumerate(self.labels):
+            labels["{}".format(i)] = F.tensor(label)
+        save_graphs(str(self.graph_path), self.graphs, labels)
+        if self.verbose:
+            print("Done saving data into cached files.")
+
+    def load(self):
+        self.graphs, labels = load_graphs(str(self.graph_path))
+
+        self.labels = []
+        for i in range(len(labels)):
+            self.labels.append(F.asnumpy(labels["{}".format(i)]))
+        # load pmpds under self.raw_path
+        with open("{}/pmpds.pkl".format(self.raw_path), "rb") as f:
+            self.pmpds = _pickle_load(f)
+        if self.verbose:
+            print("Done loading data into cached files.")
+        # sanity check
+        assert len(self.graphs) == len(self.pmpds)
+        assert len(self.graphs) == len(self.labels)
+
+    def __len__(self):
+        return len(self.graphs)
+
+    def __getitem__(self, i):
+        r"""Gets the idx-th sample.
+
+        Parameters
+        -----------
+        idx : int
+            The sample index.
+
+        Returns
+        -------
+        (dgl.DGLGraph, scipy.sparse.coo_matrix, int)
+            The graph, scipy sparse coo_matrix and its label.
+        """
+        if self._transform is None:
+            g = self.graphs[i]
+        else:
+            g = self._transform(self.graphs[i])
+        return (g, self.pmpds[i], self.labels[i])
+
+    @property
+    def save_name(self):
+        return self.name + "_dgl_graph"
+
+    @staticmethod
+    def collate_fn(cur):
+        graphs, pmpds, labels = zip(*cur)
+        batched_graphs = batch_graphs(graphs)
+        batched_pmpds = sp.block_diag(pmpds)
+        batched_labels = np.concatenate(labels, axis=0)
+        return batched_graphs, batched_pmpds, batched_labels
+
+
+def _normalize(mx):
+    """Row-normalize sparse matrix"""
+    rowsum = np.asarray(mx.sum(1))
+    mask = np.equal(rowsum, 0.0).flatten()
+    rowsum[mask] = np.nan
+    r_inv = np.power(rowsum, -1).flatten()
+    r_inv[mask] = 0.0
+    r_mat_inv = sp.diags(r_inv)
+    return r_mat_inv.dot(mx)
+
+
+def _encode_onehot(labels):
+    classes = list(sorted(set(labels)))
+    classes_dict = {
+        c: np.identity(len(classes))[i, :] for i, c in enumerate(classes)
+    }
+    labels_onehot = np.asarray(
+        list(map(classes_dict.get, labels)), dtype=np.int32
+    )
+    return labels_onehot

python/dgl/data/cluster.py

@@ -0,0 +1,132 @@
+""" CLUSTERDataset for inductive learning. """
+import os
+
+from .dgl_dataset import DGLBuiltinDataset
+from .utils import _get_dgl_url, load_graphs
+
+
+class CLUSTERDataset(DGLBuiltinDataset):
+    r"""CLUSTER dataset for semi-supervised clustering task.
+
+    Each graph contains 6 SBM clusters with sizes randomly selected between
+    [5, 35] and probabilities p = 0.55, q = 0.25. The graphs are of sizes 40
+    -190 nodes. Each node can take an input feature value in {0, 1, 2, ..., 6}
+    and values 1~6 correspond to classes 0~5 respectively, while value 0 means
+    that the class of the node is unknown. There is only one labeled node that
+    is randomly assigned to each community and most node features are set to 0.
+
+    Reference `<https://arxiv.org/pdf/2003.00982.pdf>`_
+
+    Statistics:
+
+    - Train examples: 10,000
+    - Valid examples: 1,000
+    - Test examples: 1,000
+    - Number of classes for each node: 6
+
+    Parameters
+    ----------
+    mode : str
+        Must be one of ('train', 'valid', 'test').
+        Default: 'train'
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset.
+        Default: False
+    verbose : bool
+        Whether to print out progress information.
+        Default: False
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of classes for each node.
+
+    Examples
+    --------
+    >>> from dgl.data import CLUSTERDataset
+    >>>
+    >>> trainset = CLUSTERDataset(mode='train')
+    >>>
+    >>> trainset.num_classes
+    6
+    >>> len(trainset)
+    10000
+    >>> trainset[0]
+    Graph(num_nodes=117, num_edges=4104,
+          ndata_schemes={'label': Scheme(shape=(), dtype=torch.int16),
+                         'feat': Scheme(shape=(), dtype=torch.int64)}
+          edata_schemes={'feat': Scheme(shape=(1,), dtype=torch.float32)})
+    """
+
+    def __init__(
+        self,
+        mode="train",
+        raw_dir=None,
+        force_reload=False,
+        verbose=False,
+        transform=None,
+    ):
+        self._url = _get_dgl_url("dataset/SBM_CLUSTER.zip")
+        self.mode = mode
+
+        super(CLUSTERDataset, self).__init__(
+            name="cluster",
+            url=self._url,
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        self.load()
+
+    def has_cache(self):
+        graph_path = os.path.join(
+            self.save_path, "CLUSTER_{}.bin".format(self.mode)
+        )
+        return os.path.exists(graph_path)
+
+    def load(self):
+        graph_path = os.path.join(
+            self.save_path, "CLUSTER_{}.bin".format(self.mode)
+        )
+        self._graphs, _ = load_graphs(graph_path)
+
+    @property
+    def num_classes(self):
+        r"""Number of classes for each node."""
+        return 6
+
+    def __len__(self):
+        r"""The number of examples in the dataset."""
+        return len(self._graphs)
+
+    def __getitem__(self, idx):
+        r"""Get the idx^th sample.
+
+        Parameters
+        ---------
+        idx : int
+            The sample index.
+
+        Returns
+        -------
+        :class:`dgl.DGLGraph`
+            graph structure, node features, node labels and edge features.
+
+            - ``ndata['feat']``: node features
+            - ``ndata['label']``: node labels
+            - ``edata['feat']``: edge features
+        """
+        if self._transform is None:
+            return self._graphs[idx]
+        else:
+            return self._transform(self._graphs[idx])

python/dgl/data/csv_dataset.py

@@ -0,0 +1,214 @@
+import os
+
+import numpy as np
+
+from .. import backend as F
+from ..base import DGLError
+from .dgl_dataset import DGLDataset
+from .utils import load_graphs, save_graphs, Subset
+
+
+class CSVDataset(DGLDataset):
+    """Dataset class that loads and parses graph data from CSV files.
+
+    This class requires the following additional packages:
+
+        - pyyaml >= 5.4.1
+        - pandas >= 1.1.5
+        - pydantic >= 1.9.0
+
+    The parsed graph and feature data will be cached for faster reloading. If
+    the source CSV files are modified, please specify ``force_reload=True``
+    to re-parse from them.
+
+    Parameters
+    ----------
+    data_path : str
+        Directory which contains 'meta.yaml' and CSV files
+    force_reload : bool, optional
+        Whether to reload the dataset. Default: False
+    verbose: bool, optional
+        Whether to print out progress information. Default: True.
+    ndata_parser : dict[str, callable] or callable, optional
+        Callable object which takes in the ``pandas.DataFrame`` object created from
+        CSV file, parses node data and returns a dictionary of parsed data. If given a
+        dictionary, the key is node type and the value is a callable object which is
+        used to parse data of corresponding node type. If given a single callable
+        object, such object is used to parse data of all node type data. Default: None.
+        If None, a default data parser is applied which load data directly and tries to
+        convert list into array.
+    edata_parser : dict[(str, str, str), callable], or callable, optional
+        Callable object which takes in the ``pandas.DataFrame`` object created from
+        CSV file, parses edge data and returns a dictionary of parsed data. If given a
+        dictionary, the key is edge type and the value is a callable object which is
+        used to parse data of corresponding edge type. If given a single callable
+        object, such object is used to parse data of all edge type data. Default: None.
+        If None, a default data parser is applied which load data directly and tries to
+        convert list into array.
+    gdata_parser : callable, optional
+        Callable object which takes in the ``pandas.DataFrame`` object created from
+        CSV file, parses graph data and returns a dictionary of parsed data. Default:
+        None. If None, a default data parser is applied which load data directly and
+        tries to convert list into array.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    graphs : :class:`dgl.DGLGraph`
+        Graphs of the dataset
+    data : dict
+        any available graph-level data such as graph-level feature, labels.
+
+    Examples
+    --------
+    Please refer to :ref:`guide-data-pipeline-loadcsv`.
+
+    """
+
+    META_YAML_NAME = "meta.yaml"
+
+    def __init__(
+        self,
+        data_path,
+        force_reload=False,
+        verbose=True,
+        ndata_parser=None,
+        edata_parser=None,
+        gdata_parser=None,
+        transform=None,
+    ):
+        from .csv_dataset_base import (
+            DefaultDataParser,
+            load_yaml_with_sanity_check,
+        )
+
+        self.graphs = None
+        self.data = None
+        self.ndata_parser = {} if ndata_parser is None else ndata_parser
+        self.edata_parser = {} if edata_parser is None else edata_parser
+        self.gdata_parser = gdata_parser
+        self.default_data_parser = DefaultDataParser()
+        meta_yaml_path = os.path.join(data_path, CSVDataset.META_YAML_NAME)
+        if not os.path.exists(meta_yaml_path):
+            raise DGLError(
+                "'{}' cannot be found under {}.".format(
+                    CSVDataset.META_YAML_NAME, data_path
+                )
+            )
+        self.meta_yaml = load_yaml_with_sanity_check(meta_yaml_path)
+        ds_name = self.meta_yaml.dataset_name
+        super().__init__(
+            ds_name,
+            raw_dir=os.path.dirname(meta_yaml_path),
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        """Parse node/edge data from CSV files and construct DGL.Graphs"""
+        from .csv_dataset_base import (
+            DGLGraphConstructor,
+            EdgeData,
+            GraphData,
+            NodeData,
+        )
+
+        meta_yaml = self.meta_yaml
+        base_dir = self.raw_dir
+        node_data = []
+        for meta_node in meta_yaml.node_data:
+            if meta_node is None:
+                continue
+            ntype = meta_node.ntype
+            data_parser = (
+                self.ndata_parser
+                if callable(self.ndata_parser)
+                else self.ndata_parser.get(ntype, self.default_data_parser)
+            )
+            ndata = NodeData.load_from_csv(
+                meta_node,
+                base_dir=base_dir,
+                separator=meta_yaml.separator,
+                data_parser=data_parser,
+            )
+            node_data.append(ndata)
+        edge_data = []
+        for meta_edge in meta_yaml.edge_data:
+            if meta_edge is None:
+                continue
+            etype = tuple(meta_edge.etype)
+            data_parser = (
+                self.edata_parser
+                if callable(self.edata_parser)
+                else self.edata_parser.get(etype, self.default_data_parser)
+            )
+            edata = EdgeData.load_from_csv(
+                meta_edge,
+                base_dir=base_dir,
+                separator=meta_yaml.separator,
+                data_parser=data_parser,
+            )
+            edge_data.append(edata)
+        graph_data = None
+        if meta_yaml.graph_data is not None:
+            meta_graph = meta_yaml.graph_data
+            data_parser = (
+                self.default_data_parser
+                if self.gdata_parser is None
+                else self.gdata_parser
+            )
+            graph_data = GraphData.load_from_csv(
+                meta_graph,
+                base_dir=base_dir,
+                separator=meta_yaml.separator,
+                data_parser=data_parser,
+            )
+        # construct graphs
+        self.graphs, self.data = DGLGraphConstructor.construct_graphs(
+            node_data, edge_data, graph_data
+        )
+        if len(self.data) == 1:
+            self.labels = list(self.data.values())[0]
+
+    def has_cache(self):
+        graph_path = os.path.join(self.save_path, self.name + ".bin")
+        if os.path.exists(graph_path):
+            return True
+
+        return False
+
+    def save(self):
+        if self.graphs is None:
+            raise DGLError("No graphs available in dataset")
+        graph_path = os.path.join(self.save_path, self.name + ".bin")
+        save_graphs(graph_path, self.graphs, labels=self.data)
+
+    def load(self):
+        graph_path = os.path.join(self.save_path, self.name + ".bin")
+        self.graphs, self.data = load_graphs(graph_path)
+        if len(self.data) == 1:
+            self.labels = list(self.data.values())[0]
+
+    def __getitem__(self, i):
+        if F.is_tensor(i) and F.ndim(i) == 1:
+            return Subset(self, F.copy_to(i, F.cpu()))
+
+        if self._transform is None:
+            g = self.graphs[i]
+        else:
+            g = self._transform(self.graphs[i])
+
+        if len(self.data) == 1:
+            return g, self.labels[i]
+        elif len(self.data) > 0:
+            data = {k: v[i] for (k, v) in self.data.items()}
+            return g, data
+        else:
+            return g
+
+    def __len__(self):
+        return len(self.graphs)

python/dgl/data/csv_dataset_base.py

@@ -0,0 +1,386 @@
+import ast
+import os
+from typing import Callable, List, Optional
+
+import numpy as np
+import pandas as pd
+import pydantic as dt
+import yaml
+
+from .. import backend as F
+from ..base import dgl_warning, DGLError
+from ..convert import heterograph as dgl_heterograph
+
+
+class MetaNode(dt.BaseModel):
+    """Class of node_data in YAML. Internal use only."""
+
+    file_name: str
+    ntype: Optional[str] = "_V"
+    graph_id_field: Optional[str] = "graph_id"
+    node_id_field: Optional[str] = "node_id"
+
+
+class MetaEdge(dt.BaseModel):
+    """Class of edge_data in YAML. Internal use only."""
+
+    file_name: str
+    etype: Optional[List[str]] = ["_V", "_E", "_V"]
+    graph_id_field: Optional[str] = "graph_id"
+    src_id_field: Optional[str] = "src_id"
+    dst_id_field: Optional[str] = "dst_id"
+
+
+class MetaGraph(dt.BaseModel):
+    """Class of graph_data in YAML. Internal use only."""
+
+    file_name: str
+    graph_id_field: Optional[str] = "graph_id"
+
+
+class MetaYaml(dt.BaseModel):
+    """Class of YAML. Internal use only."""
+
+    version: Optional[str] = "1.0.0"
+    dataset_name: str
+    separator: Optional[str] = ","
+    node_data: List[MetaNode]
+    edge_data: List[MetaEdge]
+    graph_data: Optional[MetaGraph] = None
+
+
+def load_yaml_with_sanity_check(yaml_file):
+    """Load yaml and do sanity check. Internal use only."""
+    with open(yaml_file) as f:
+        yaml_data = yaml.load(f, Loader=yaml.loader.SafeLoader)
+        try:
+            meta_yaml = MetaYaml(**yaml_data)
+        except dt.ValidationError as e:
+            print("Details of pydantic.ValidationError:\n{}".format(e.json()))
+            raise DGLError(
+                "Validation Error for YAML fields. Details are shown above."
+            )
+        if meta_yaml.version != "1.0.0":
+            raise DGLError(
+                "Invalid CSVDataset version {}. Supported versions: '1.0.0'".format(
+                    meta_yaml.version
+                )
+            )
+        ntypes = [meta.ntype for meta in meta_yaml.node_data]
+        if len(ntypes) > len(set(ntypes)):
+            raise DGLError(
+                "Each node CSV file must have a unique node type name, but found duplicate node type: {}.".format(
+                    ntypes
+                )
+            )
+        etypes = [tuple(meta.etype) for meta in meta_yaml.edge_data]
+        if len(etypes) > len(set(etypes)):
+            raise DGLError(
+                "Each edge CSV file must have a unique edge type name, but found duplicate edge type: {}.".format(
+                    etypes
+                )
+            )
+        return meta_yaml
+
+
+def _validate_data_length(data_dict):
+    len_dict = {k: len(v) for k, v in data_dict.items()}
+    lst = list(len_dict.values())
+    res = lst.count(lst[0]) == len(lst)
+    if not res:
+        raise DGLError(
+            "All data are required to have same length while some of them does not. Length of data={}".format(
+                str(len_dict)
+            )
+        )
+
+
+def _tensor(data, dtype=None):
+    """Float32 is the default dtype for float tensor in DGL
+    so let's cast float64 into float32 to avoid dtype mismatch.
+    """
+    ret = F.tensor(data, dtype)
+    if F.dtype(ret) == F.float64:
+        ret = F.tensor(ret, dtype=F.float32)
+    return ret
+
+
+class BaseData:
+    """Class of base data which is inherited by Node/Edge/GraphData. Internal use only."""
+
+    @staticmethod
+    def read_csv(file_name, base_dir, separator):
+        csv_path = file_name
+        if base_dir is not None:
+            csv_path = os.path.join(base_dir, csv_path)
+        return pd.read_csv(csv_path, sep=separator)
+
+    @staticmethod
+    def pop_from_dataframe(df: pd.DataFrame, item: str):
+        ret = None
+        try:
+            ret = df.pop(item).to_numpy().squeeze()
+        except KeyError:
+            pass
+        return ret
+
+
+class NodeData(BaseData):
+    """Class of node data which is used for DGLGraph construction. Internal use only."""
+
+    def __init__(self, node_id, data, type=None, graph_id=None):
+        self.id = np.array(node_id)
+        self.data = data
+        self.type = type if type is not None else "_V"
+        self.graph_id = (
+            np.array(graph_id)
+            if graph_id is not None
+            else np.full(len(node_id), 0)
+        )
+        _validate_data_length(
+            {**{"id": self.id, "graph_id": self.graph_id}, **self.data}
+        )
+
+    @staticmethod
+    def load_from_csv(
+        meta: MetaNode, data_parser: Callable, base_dir=None, separator=","
+    ):
+        df = BaseData.read_csv(meta.file_name, base_dir, separator)
+        node_ids = BaseData.pop_from_dataframe(df, meta.node_id_field)
+        graph_ids = BaseData.pop_from_dataframe(df, meta.graph_id_field)
+        if node_ids is None:
+            raise DGLError(
+                "Missing node id field [{}] in file [{}].".format(
+                    meta.node_id_field, meta.file_name
+                )
+            )
+        ntype = meta.ntype
+        ndata = data_parser(df)
+        return NodeData(node_ids, ndata, type=ntype, graph_id=graph_ids)
+
+    @staticmethod
+    def to_dict(node_data: List["NodeData"]) -> dict:
+        # node_ids could be numeric or non-numeric values, but duplication is not allowed.
+        node_dict = {}
+        for n_data in node_data:
+            graph_ids = np.unique(n_data.graph_id)
+            for graph_id in graph_ids:
+                idx = n_data.graph_id == graph_id
+                ids = n_data.id[idx]
+                u_ids, u_indices, u_counts = np.unique(
+                    ids, return_index=True, return_counts=True
+                )
+                if len(ids) > len(u_ids):
+                    raise DGLError(
+                        "Node IDs are required to be unique but the following ids are duplicate: {}".format(
+                            u_ids[u_counts > 1]
+                        )
+                    )
+                if graph_id not in node_dict:
+                    node_dict[graph_id] = {}
+                node_dict[graph_id][n_data.type] = {
+                    "mapping": {
+                        index: i for i, index in enumerate(ids[u_indices])
+                    },
+                    "data": {
+                        k: _tensor(v[idx][u_indices])
+                        for k, v in n_data.data.items()
+                    },
+                    "dtype": ids.dtype,
+                }
+        return node_dict
+
+
+class EdgeData(BaseData):
+    """Class of edge data which is used for DGLGraph construction. Internal use only."""
+
+    def __init__(self, src_id, dst_id, data, type=None, graph_id=None):
+        self.src = np.array(src_id)
+        self.dst = np.array(dst_id)
+        self.data = data
+        self.type = type if type is not None else ("_V", "_E", "_V")
+        self.graph_id = (
+            np.array(graph_id)
+            if graph_id is not None
+            else np.full(len(src_id), 0)
+        )
+        _validate_data_length(
+            {
+                **{"src": self.src, "dst": self.dst, "graph_id": self.graph_id},
+                **self.data,
+            }
+        )
+
+    @staticmethod
+    def load_from_csv(
+        meta: MetaEdge, data_parser: Callable, base_dir=None, separator=","
+    ):
+        df = BaseData.read_csv(meta.file_name, base_dir, separator)
+        src_ids = BaseData.pop_from_dataframe(df, meta.src_id_field)
+        if src_ids is None:
+            raise DGLError(
+                "Missing src id field [{}] in file [{}].".format(
+                    meta.src_id_field, meta.file_name
+                )
+            )
+        dst_ids = BaseData.pop_from_dataframe(df, meta.dst_id_field)
+        if dst_ids is None:
+            raise DGLError(
+                "Missing dst id field [{}] in file [{}].".format(
+                    meta.dst_id_field, meta.file_name
+                )
+            )
+        graph_ids = BaseData.pop_from_dataframe(df, meta.graph_id_field)
+        etype = tuple(meta.etype)
+        edata = data_parser(df)
+        return EdgeData(src_ids, dst_ids, edata, type=etype, graph_id=graph_ids)
+
+    @staticmethod
+    def to_dict(edge_data: List["EdgeData"], node_dict: dict) -> dict:
+        edge_dict = {}
+        for e_data in edge_data:
+            (src_type, e_type, dst_type) = e_data.type
+            graph_ids = np.unique(e_data.graph_id)
+            for graph_id in graph_ids:
+                if graph_id in edge_dict and e_data.type in edge_dict[graph_id]:
+                    raise DGLError(
+                        f"Duplicate edge type[{e_data.type}] for same graph[{graph_id}], please place the same edge_type for same graph into single EdgeData."
+                    )
+                idx = e_data.graph_id == graph_id
+                src_mapping = node_dict[graph_id][src_type]["mapping"]
+                dst_mapping = node_dict[graph_id][dst_type]["mapping"]
+                orig_src_ids = e_data.src[idx].astype(
+                    node_dict[graph_id][src_type]["dtype"]
+                )
+                orig_dst_ids = e_data.dst[idx].astype(
+                    node_dict[graph_id][dst_type]["dtype"]
+                )
+                src_ids = [src_mapping[index] for index in orig_src_ids]
+                dst_ids = [dst_mapping[index] for index in orig_dst_ids]
+                if graph_id not in edge_dict:
+                    edge_dict[graph_id] = {}
+                edge_dict[graph_id][e_data.type] = {
+                    "edges": (_tensor(src_ids), _tensor(dst_ids)),
+                    "data": {
+                        k: _tensor(v[idx]) for k, v in e_data.data.items()
+                    },
+                }
+        return edge_dict
+
+
+class GraphData(BaseData):
+    """Class of graph data which is used for DGLGraph construction. Internal use only."""
+
+    def __init__(self, graph_id, data):
+        self.graph_id = np.array(graph_id)
+        self.data = data
+        _validate_data_length({**{"graph_id": self.graph_id}, **self.data})
+
+    @staticmethod
+    def load_from_csv(
+        meta: MetaGraph, data_parser: Callable, base_dir=None, separator=","
+    ):
+        df = BaseData.read_csv(meta.file_name, base_dir, separator)
+        graph_ids = BaseData.pop_from_dataframe(df, meta.graph_id_field)
+        if graph_ids is None:
+            raise DGLError(
+                "Missing graph id field [{}] in file [{}].".format(
+                    meta.graph_id_field, meta.file_name
+                )
+            )
+        gdata = data_parser(df)
+        return GraphData(graph_ids, gdata)
+
+    @staticmethod
+    def to_dict(graph_data: "GraphData", graphs_dict: dict) -> dict:
+        missing_ids = np.setdiff1d(
+            np.array(list(graphs_dict.keys())), graph_data.graph_id
+        )
+        if len(missing_ids) > 0:
+            raise DGLError(
+                "Found following graph ids in node/edge CSVs but not in graph CSV: {}.".format(
+                    missing_ids
+                )
+            )
+        graph_ids = graph_data.graph_id
+        graphs = []
+        for graph_id in graph_ids:
+            if graph_id not in graphs_dict:
+                graphs_dict[graph_id] = dgl_heterograph(
+                    {("_V", "_E", "_V"): ([], [])}
+                )
+        for graph_id in graph_ids:
+            graphs.append(graphs_dict[graph_id])
+        data = {
+            k: F.reshape(_tensor(v), (len(graphs), -1))
+            for k, v in graph_data.data.items()
+        }
+        return graphs, data
+
+
+class DGLGraphConstructor:
+    """Class for constructing DGLGraph from Node/Edge/Graph data. Internal use only."""
+
+    @staticmethod
+    def construct_graphs(node_data, edge_data, graph_data=None):
+        if not isinstance(node_data, list):
+            node_data = [node_data]
+        if not isinstance(edge_data, list):
+            edge_data = [edge_data]
+        node_dict = NodeData.to_dict(node_data)
+        edge_dict = EdgeData.to_dict(edge_data, node_dict)
+        graph_dict = DGLGraphConstructor._construct_graphs(node_dict, edge_dict)
+        if graph_data is None:
+            graph_data = GraphData(np.full(1, 0), {})
+        graphs, data = GraphData.to_dict(graph_data, graph_dict)
+        return graphs, data
+
+    @staticmethod
+    def _construct_graphs(node_dict, edge_dict):
+        graph_dict = {}
+        for graph_id in node_dict:
+            if graph_id not in edge_dict:
+                edge_dict[graph_id][("_V", "_E", "_V")] = {"edges": ([], [])}
+            graph = dgl_heterograph(
+                {
+                    etype: edata["edges"]
+                    for etype, edata in edge_dict[graph_id].items()
+                },
+                num_nodes_dict={
+                    ntype: len(ndata["mapping"])
+                    for ntype, ndata in node_dict[graph_id].items()
+                },
+            )
+
+            def assign_data(type, src_data, dst_data):
+                for key, value in src_data.items():
+                    dst_data[type].data[key] = value
+
+            for type, data in node_dict[graph_id].items():
+                assign_data(type, data["data"], graph.nodes)
+            for (type), data in edge_dict[graph_id].items():
+                assign_data(type, data["data"], graph.edges)
+            graph_dict[graph_id] = graph
+        return graph_dict
+
+
+class DefaultDataParser:
+    """Default data parser for CSVDataset. It
+    1. ignores any columns which does not have a header.
+    2. tries to convert to list of numeric values(generated by
+        np.array().tolist()) if cell data is a str separated by ','.
+    3. read data and infer data type directly, otherwise.
+    """
+
+    def __call__(self, df: pd.DataFrame):
+        data = {}
+        for header in df:
+            if "Unnamed" in header:
+                dgl_warning("Unnamed column is found. Ignored...")
+                continue
+            dt = df[header].to_numpy().squeeze()
+            if len(dt) > 0 and isinstance(dt[0], str):
+                # probably consists of list of numeric values
+                dt = np.array([ast.literal_eval(row) for row in dt])
+            data[header] = dt
+        return data

python/dgl/data/dgl_dataset.py

@@ -0,0 +1,349 @@
+"""Basic DGL Dataset
+"""
+
+from __future__ import absolute_import
+
+import abc
+import hashlib
+import os
+import traceback
+
+from ..utils import retry_method_with_fix
+from .utils import download, extract_archive, get_download_dir, makedirs
+
+
+class DGLDataset(object):
+    r"""The basic DGL dataset for creating graph datasets.
+    This class defines a basic template class for DGL Dataset.
+    The following steps will be executed automatically:
+
+      1. Check whether there is a dataset cache on disk
+         (already processed and stored on the disk) by
+         invoking ``has_cache()``. If true, goto 5.
+      2. Call ``download()`` to download the data if ``url`` is not None.
+      3. Call ``process()`` to process the data.
+      4. Call ``save()`` to save the processed dataset on disk and goto 6.
+      5. Call ``load()`` to load the processed dataset from disk.
+      6. Done.
+
+    Users can overwite these functions with their
+    own data processing logic.
+
+    Parameters
+    ----------
+    name : str
+        Name of the dataset
+    url : str
+        Url to download the raw dataset. Default: None
+    raw_dir : str
+        Specifying the directory that will store the
+        downloaded data or the directory that
+        already stores the input data.
+        Default: ~/.dgl/
+    save_dir : str
+        Directory to save the processed dataset.
+        Default: same as raw_dir
+    hash_key : tuple
+        A tuple of values as the input for the hash function.
+        Users can distinguish instances (and their caches on the disk)
+        from the same dataset class by comparing the hash values.
+        Default: (), the corresponding hash value is ``'f9065fa7'``.
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    url : str
+        The URL to download the dataset
+    name : str
+        The dataset name
+    raw_dir : str
+        Directory to store all the downloaded raw datasets.
+    raw_path : str
+        Path to the downloaded raw dataset folder. An alias for
+        ``os.path.join(self.raw_dir, self.name)``.
+    save_dir : str
+        Directory to save all the processed datasets.
+    save_path : str
+        Path to the processed dataset folder. An alias for
+        ``os.path.join(self.save_dir, self.name)``.
+    verbose : bool
+        Whether to print more runtime information.
+    hash : str
+        Hash value for the dataset and the setting.
+    """
+
+    def __init__(
+        self,
+        name,
+        url=None,
+        raw_dir=None,
+        save_dir=None,
+        hash_key=(),
+        force_reload=False,
+        verbose=False,
+        transform=None,
+    ):
+        self._name = name
+        self._url = url
+        self._force_reload = force_reload
+        self._verbose = verbose
+        self._hash_key = hash_key
+        self._hash = self._get_hash()
+        self._transform = transform
+
+        # if no dir is provided, the default dgl download dir is used.
+        if raw_dir is None:
+            self._raw_dir = get_download_dir()
+        else:
+            self._raw_dir = raw_dir
+
+        if save_dir is None:
+            self._save_dir = self._raw_dir
+        else:
+            self._save_dir = save_dir
+
+        self._load()
+
+    def download(self):
+        r"""Overwite to realize your own logic of downloading data.
+
+        It is recommended to download the to the :obj:`self.raw_dir`
+        folder. Can be ignored if the dataset is
+        already in :obj:`self.raw_dir`.
+        """
+        pass
+
+    def save(self):
+        r"""Overwite to realize your own logic of
+        saving the processed dataset into files.
+
+        It is recommended to use ``dgl.data.utils.save_graphs``
+        to save dgl graph into files and use
+        ``dgl.data.utils.save_info`` to save extra
+        information into files.
+        """
+        pass
+
+    def load(self):
+        r"""Overwite to realize your own logic of
+        loading the saved dataset from files.
+
+        It is recommended to use ``dgl.data.utils.load_graphs``
+        to load dgl graph from files and use
+        ``dgl.data.utils.load_info`` to load extra information
+        into python dict object.
+        """
+        pass
+
+    @abc.abstractmethod
+    def process(self):
+        r"""Overwrite to realize your own logic of processing the input data."""
+        pass
+
+    def has_cache(self):
+        r"""Overwrite to realize your own logic of
+        deciding whether there exists a cached dataset.
+
+        By default False.
+        """
+        return False
+
+    @retry_method_with_fix(download)
+    def _download(self):
+        """Download dataset by calling ``self.download()``
+        if the dataset does not exists under ``self.raw_path``.
+
+        By default ``self.raw_path = os.path.join(self.raw_dir, self.name)``
+        One can overwrite ``raw_path()`` function to change the path.
+        """
+        if os.path.exists(self.raw_path):  # pragma: no cover
+            return
+
+        makedirs(self.raw_dir)
+        self.download()
+
+    def _load(self):
+        """Entry point from __init__ to load the dataset.
+
+        If cache exists:
+
+          - Load the dataset from saved dgl graph and information files.
+          - If loadin process fails, re-download and process the dataset.
+
+        else:
+
+          - Download the dataset if needed.
+          - Process the dataset and build the dgl graph.
+          - Save the processed dataset into files.
+        """
+        load_flag = not self._force_reload and self.has_cache()
+
+        if load_flag:
+            try:
+                self.load()
+                if self.verbose:
+                    print("Done loading data from cached files.")
+            except KeyboardInterrupt:
+                raise
+            except:
+                load_flag = False
+                if self.verbose:
+                    print(traceback.format_exc())
+                    print("Loading from cache failed, re-processing.")
+
+        if not load_flag:
+            self._download()
+            self.process()
+            self.save()
+            if self.verbose:
+                print("Done saving data into cached files.")
+
+    def _get_hash(self):
+        """Compute the hash of the input tuple
+
+        Example
+        -------
+        Assume `self._hash_key = (10, False, True)`
+
+        >>> hash_value = self._get_hash()
+        >>> hash_value
+        'a770b222'
+        """
+        hash_func = hashlib.sha1()
+        hash_func.update(str(self._hash_key).encode("utf-8"))
+        return hash_func.hexdigest()[:8]
+
+    def _get_hash_url_suffix(self):
+        """Get the suffix based on the hash value of the url."""
+        if self._url is None:
+            return ""
+        else:
+            hash_func = hashlib.sha1()
+            hash_func.update(str(self._url).encode("utf-8"))
+            return "_" + hash_func.hexdigest()[:8]
+
+    @property
+    def url(self):
+        r"""Get url to download the raw dataset."""
+        return self._url
+
+    @property
+    def name(self):
+        r"""Name of the dataset."""
+        return self._name
+
+    @property
+    def raw_dir(self):
+        r"""Raw file directory contains the input data folder."""
+        return self._raw_dir
+
+    @property
+    def raw_path(self):
+        r"""Directory contains the input data files.
+        By default raw_path = os.path.join(self.raw_dir, self.name)
+        """
+        return os.path.join(
+            self.raw_dir, self.name + self._get_hash_url_suffix()
+        )
+
+    @property
+    def save_dir(self):
+        r"""Directory to save the processed dataset."""
+        return self._save_dir
+
+    @property
+    def save_path(self):
+        r"""Path to save the processed dataset."""
+        return os.path.join(
+            self.save_dir, self.name + self._get_hash_url_suffix()
+        )
+
+    @property
+    def verbose(self):
+        r"""Whether to print information."""
+        return self._verbose
+
+    @property
+    def hash(self):
+        r"""Hash value for the dataset and the setting."""
+        return self._hash
+
+    @abc.abstractmethod
+    def __getitem__(self, idx):
+        r"""Gets the data object at index."""
+        pass
+
+    @abc.abstractmethod
+    def __len__(self):
+        r"""The number of examples in the dataset."""
+        pass
+
+    def __repr__(self):
+        return (
+            f'Dataset("{self.name}", num_graphs={len(self)},'
+            + f" save_path={self.save_path})"
+        )
+
+
+class DGLBuiltinDataset(DGLDataset):
+    r"""The Basic DGL Builtin Dataset.
+
+    Parameters
+    ----------
+    name : str
+        Name of the dataset.
+    url : str
+        Url to download the raw dataset.
+    raw_dir : str
+        Specifying the directory that will store the
+        downloaded data or the directory that
+        already stores the input data.
+        Default: ~/.dgl/
+    hash_key : tuple
+        A tuple of values as the input for the hash function.
+        Users can distinguish instances (and their caches on the disk)
+        from the same dataset class by comparing the hash values.
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: False
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+    """
+
+    def __init__(
+        self,
+        name,
+        url,
+        raw_dir=None,
+        hash_key=(),
+        force_reload=False,
+        verbose=False,
+        transform=None,
+    ):
+        super(DGLBuiltinDataset, self).__init__(
+            name,
+            url=url,
+            raw_dir=raw_dir,
+            save_dir=None,
+            hash_key=hash_key,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def download(self):
+        r"""Automatically download data and extract it."""
+        if self.url is not None:
+            zip_file_path = os.path.join(self.raw_dir, self.name + ".zip")
+            download(self.url, path=zip_file_path)
+            extract_archive(zip_file_path, self.raw_path)

python/dgl/data/fakenews.py

@@ -0,0 +1,255 @@
+import os
+
+import numpy as np
+import scipy.sparse as sp
+
+from .. import backend as F
+from ..convert import graph
+from .dgl_dataset import DGLBuiltinDataset
+from .utils import _get_dgl_url, load_graphs, load_info, save_graphs, save_info
+
+
+class FakeNewsDataset(DGLBuiltinDataset):
+    r"""Fake News Graph Classification dataset.
+
+    The dataset is composed of two sets of tree-structured fake/real
+    news propagation graphs extracted from Twitter. Different from
+    most of the benchmark datasets for the graph classification task,
+    the graphs in this dataset are directed tree-structured graphs where
+    the root node represents the news, the leaf nodes are Twitter users
+    who retweeted the root news. Besides, the node features are encoded
+    user historical tweets using different pretrained language models:
+
+    - bert: the 768-dimensional node feature composed of Twitter user historical tweets encoded by the bert-as-service
+    - content: the 310-dimensional node feature composed of a 300-dimensional “spacy” vector plus a 10-dimensional “profile” vector
+    - profile: the 10-dimensional node feature composed of ten Twitter user profile attributes.
+    - spacy: the 300-dimensional node feature composed of Twitter user historical tweets encoded by the spaCy word2vec encoder.
+
+    Reference: <https://github.com/safe-graph/GNN-FakeNews>
+
+    Note: this dataset is for academic use only, and commercial use is prohibited.
+
+    Statistics:
+
+        Politifact:
+
+        - Graphs: 314
+        - Nodes: 41,054
+        - Edges: 40,740
+        - Classes:
+
+            - Fake: 157
+            - Real: 157
+
+        - Node feature size:
+
+            - bert: 768
+            - content: 310
+            - profile: 10
+            - spacy: 300
+
+        Gossipcop:
+
+        - Graphs: 5,464
+        - Nodes: 314,262
+        - Edges: 308,798
+        - Classes:
+
+            - Fake: 2,732
+            - Real: 2,732
+
+        - Node feature size:
+
+            - bert: 768
+            - content: 310
+            - profile: 10
+            - spacy: 300
+
+    Parameters
+    ----------
+    name : str
+        Name of the dataset (gossipcop, or politifact)
+    feature_name : str
+        Name of the feature (bert, content, profile, or spacy)
+    raw_dir : str
+        Specifying the directory that will store the
+        downloaded data or the directory that
+        already stores the input data.
+        Default: ~/.dgl/
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    name : str
+        Name of the dataset (gossipcop, or politifact)
+    num_classes : int
+        Number of label classes
+    num_graphs : int
+        Number of graphs
+    graphs : list
+        A list of DGLGraph objects
+    labels : Tensor
+        Graph labels
+    feature_name : str
+        Name of the feature (bert, content, profile, or spacy)
+    feature : Tensor
+        Node features
+    train_mask : Tensor
+        Mask of training set
+    val_mask : Tensor
+        Mask of validation set
+    test_mask : Tensor
+        Mask of testing set
+
+    Examples
+    --------
+    >>> dataset = FakeNewsDataset('gossipcop', 'bert')
+    >>> graph, label = dataset[0]
+    >>> num_classes = dataset.num_classes
+    >>> feat = dataset.feature
+    >>> labels = dataset.labels
+    """
+    file_urls = {
+        "gossipcop": "dataset/FakeNewsGOS.zip",
+        "politifact": "dataset/FakeNewsPOL.zip",
+    }
+
+    def __init__(self, name, feature_name, raw_dir=None, transform=None):
+        assert name in [
+            "gossipcop",
+            "politifact",
+        ], "Only supports 'gossipcop' or 'politifact'."
+        url = _get_dgl_url(self.file_urls[name])
+
+        assert feature_name in [
+            "bert",
+            "content",
+            "profile",
+            "spacy",
+        ], "Only supports 'bert', 'content', 'profile', or 'spacy'"
+        self.feature_name = feature_name
+        super(FakeNewsDataset, self).__init__(
+            name=name, url=url, raw_dir=raw_dir, transform=transform
+        )
+
+    def process(self):
+        """process raw data to graph, labels and masks"""
+        self.labels = F.tensor(
+            np.load(os.path.join(self.raw_path, "graph_labels.npy"))
+        )
+        num_graphs = self.labels.shape[0]
+
+        node_graph_id = np.load(
+            os.path.join(self.raw_path, "node_graph_id.npy")
+        )
+        edges = np.genfromtxt(
+            os.path.join(self.raw_path, "A.txt"), delimiter=",", dtype=int
+        )
+        src = edges[:, 0]
+        dst = edges[:, 1]
+        g = graph((src, dst))
+
+        node_idx_list = []
+        for idx in range(np.max(node_graph_id) + 1):
+            node_idx = np.where(node_graph_id == idx)
+            node_idx_list.append(node_idx[0])
+
+        self.graphs = [g.subgraph(node_idx) for node_idx in node_idx_list]
+
+        train_idx = np.load(os.path.join(self.raw_path, "train_idx.npy"))
+        val_idx = np.load(os.path.join(self.raw_path, "val_idx.npy"))
+        test_idx = np.load(os.path.join(self.raw_path, "test_idx.npy"))
+        train_mask = np.zeros(num_graphs, dtype=np.bool_)
+        val_mask = np.zeros(num_graphs, dtype=np.bool_)
+        test_mask = np.zeros(num_graphs, dtype=np.bool_)
+        train_mask[train_idx] = True
+        val_mask[val_idx] = True
+        test_mask[test_idx] = True
+        self.train_mask = F.tensor(train_mask)
+        self.val_mask = F.tensor(val_mask)
+        self.test_mask = F.tensor(test_mask)
+
+        feature_file = "new_" + self.feature_name + "_feature.npz"
+        self.feature = F.tensor(
+            sp.load_npz(os.path.join(self.raw_path, feature_file)).todense()
+        )
+
+    def save(self):
+        """save the graph list and the labels"""
+        save_graphs(str(self.graph_path), self.graphs)
+        save_info(
+            self.info_path,
+            {
+                "label": self.labels,
+                "feature": self.feature,
+                "train_mask": self.train_mask,
+                "val_mask": self.val_mask,
+                "test_mask": self.test_mask,
+            },
+        )
+
+    @property
+    def graph_path(self):
+        return os.path.join(self.save_path, self.name + "_dgl_graph.bin")
+
+    @property
+    def info_path(self):
+        return os.path.join(self.save_path, self.name + "_dgl_graph.pkl")
+
+    def has_cache(self):
+        """check whether there are processed data in `self.save_path`"""
+        return os.path.exists(self.graph_path) and os.path.exists(
+            self.info_path
+        )
+
+    def load(self):
+        """load processed data from directory `self.save_path`"""
+        graphs, _ = load_graphs(str(self.graph_path))
+        info = load_info(str(self.info_path))
+        self.graphs = graphs
+        self.labels = info["label"]
+        self.feature = info["feature"]
+
+        self.train_mask = info["train_mask"]
+        self.val_mask = info["val_mask"]
+        self.test_mask = info["test_mask"]
+
+    @property
+    def num_classes(self):
+        """Number of classes for each graph, i.e. number of prediction tasks."""
+        return 2
+
+    @property
+    def num_graphs(self):
+        """Number of graphs."""
+        return self.labels.shape[0]
+
+    def __getitem__(self, i):
+        r"""Get graph and label by index
+
+        Parameters
+        ----------
+        i : int
+            Item index
+
+        Returns
+        -------
+        (:class:`dgl.DGLGraph`, Tensor)
+        """
+        if self._transform is None:
+            g = self.graphs[i]
+        else:
+            g = self._transform(self.graphs[i])
+        return g, self.labels[i]
+
+    def __len__(self):
+        r"""Number of graphs in the dataset.
+
+        Return
+        -------
+        int
+        """
+        return len(self.graphs)

python/dgl/data/flickr.py

@@ -0,0 +1,178 @@
+"""Flickr Dataset"""
+import json
+import os
+
+import numpy as np
+import scipy.sparse as sp
+
+from .. import backend as F
+from ..convert import from_scipy
+from ..transforms import reorder_graph
+from .dgl_dataset import DGLBuiltinDataset
+from .utils import _get_dgl_url, generate_mask_tensor, load_graphs, save_graphs
+
+
+class FlickrDataset(DGLBuiltinDataset):
+    r"""Flickr dataset for node classification from `GraphSAINT: Graph Sampling Based Inductive
+    Learning Method <https://arxiv.org/abs/1907.04931>`_
+
+    The task of this dataset is categorizing types of images based on the descriptions and common
+    properties of online images.
+
+    Flickr dataset statistics:
+
+    - Nodes: 89,250
+    - Edges: 899,756
+    - Number of classes: 7
+    - Node feature size: 500
+
+    Parameters
+    ----------
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset.
+        Default: False
+    verbose : bool
+        Whether to print out progress information.
+        Default: False
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+    reorder : bool
+        Whether to reorder the graph using :func:`~dgl.reorder_graph`.
+        Default: False.
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of node classes
+
+    Examples
+    --------
+    >>> from dgl.data import FlickrDataset
+    >>> dataset = FlickrDataset()
+    >>> dataset.num_classes
+    7
+    >>> g = dataset[0]
+    >>> # get node feature
+    >>> feat = g.ndata['feat']
+    >>> # get node labels
+    >>> labels = g.ndata['label']
+    >>> # get data split
+    >>> train_mask = g.ndata['train_mask']
+    >>> val_mask = g.ndata['val_mask']
+    >>> test_mask = g.ndata['test_mask']
+    """
+
+    def __init__(
+        self,
+        raw_dir=None,
+        force_reload=False,
+        verbose=False,
+        transform=None,
+        reorder=False,
+    ):
+        _url = _get_dgl_url("dataset/flickr.zip")
+        self._reorder = reorder
+        super(FlickrDataset, self).__init__(
+            name="flickr",
+            raw_dir=raw_dir,
+            url=_url,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        """process raw data to graph, labels and masks"""
+        coo_adj = sp.load_npz(os.path.join(self.raw_path, "adj_full.npz"))
+        g = from_scipy(coo_adj)
+
+        features = np.load(os.path.join(self.raw_path, "feats.npy"))
+        features = F.tensor(features, dtype=F.float32)
+
+        y = [-1] * features.shape[0]
+        with open(os.path.join(self.raw_path, "class_map.json")) as f:
+            class_map = json.load(f)
+            for key, item in class_map.items():
+                y[int(key)] = item
+        labels = F.tensor(np.array(y), dtype=F.int64)
+
+        with open(os.path.join(self.raw_path, "role.json")) as f:
+            role = json.load(f)
+
+        train_mask = np.zeros(features.shape[0], dtype=bool)
+        train_mask[role["tr"]] = True
+
+        val_mask = np.zeros(features.shape[0], dtype=bool)
+        val_mask[role["va"]] = True
+
+        test_mask = np.zeros(features.shape[0], dtype=bool)
+        test_mask[role["te"]] = True
+
+        g.ndata["feat"] = features
+        g.ndata["label"] = labels
+        g.ndata["train_mask"] = generate_mask_tensor(train_mask)
+        g.ndata["val_mask"] = generate_mask_tensor(val_mask)
+        g.ndata["test_mask"] = generate_mask_tensor(test_mask)
+
+        if self._reorder:
+            self._graph = reorder_graph(
+                g,
+                node_permute_algo="rcmk",
+                edge_permute_algo="dst",
+                store_ids=False,
+            )
+        else:
+            self._graph = g
+
+    def has_cache(self):
+        graph_path = os.path.join(self.save_path, "dgl_graph.bin")
+        return os.path.exists(graph_path)
+
+    def save(self):
+        graph_path = os.path.join(self.save_path, "dgl_graph.bin")
+        save_graphs(graph_path, self._graph)
+
+    def load(self):
+        graph_path = os.path.join(self.save_path, "dgl_graph.bin")
+        g, _ = load_graphs(graph_path)
+        self._graph = g[0]
+
+    @property
+    def num_classes(self):
+        return 7
+
+    def __len__(self):
+        r"""The number of graphs in the dataset."""
+        return 1
+
+    def __getitem__(self, idx):
+        r"""Get graph object
+
+        Parameters
+        ----------
+        idx : int
+            Item index, FlickrDataset has only one graph object
+
+        Returns
+        -------
+        :class:`dgl.DGLGraph`
+
+            The graph contains:
+
+            - ``ndata['label']``: node label
+            - ``ndata['feat']``: node feature
+            - ``ndata['train_mask']``: mask for training node set
+            - ``ndata['val_mask']``: mask for validation node set
+            - ``ndata['test_mask']``: mask for test node set
+
+        """
+        assert idx == 0, "This dataset has only one graph"
+        if self._transform is None:
+            return self._graph
+        else:
+            return self._transform(self._graph)

python/dgl/data/fraud.py

@@ -0,0 +1,415 @@
+"""Fraud Dataset
+"""
+import os
+
+import numpy as np
+from scipy import io
+
+from .. import backend as F
+from ..convert import heterograph
+from .dgl_dataset import DGLBuiltinDataset
+from .utils import _get_dgl_url, load_graphs, save_graphs
+
+
+class FraudDataset(DGLBuiltinDataset):
+    r"""Fraud node prediction dataset.
+
+    The dataset includes two multi-relational graphs extracted from Yelp and Amazon
+    where nodes represent fraudulent reviews or fraudulent reviewers.
+
+    It was first proposed in a CIKM'20 paper <https://arxiv.org/pdf/2008.08692.pdf> and
+    has been used by a recent WWW'21 paper <https://ponderly.github.io/pub/PCGNN_WWW2021.pdf>
+    as a benchmark. Another paper <https://arxiv.org/pdf/2104.01404.pdf> also takes
+    the dataset as an example to study the non-homophilous graphs. This dataset is built
+    upon industrial data and has rich relational information and unique properties like
+    class-imbalance and feature inconsistency, which makes the dataset be a good instance
+    to investigate how GNNs perform on real-world noisy graphs. These graphs are bidirected
+    and not self connected.
+
+    Reference: <https://github.com/YingtongDou/CARE-GNN>
+
+    Parameters
+    ----------
+    name : str
+        Name of the dataset
+    raw_dir : str
+        Specifying the directory that will store the
+        downloaded data or the directory that
+        already stores the input data.
+        Default: ~/.dgl/
+    random_seed : int
+        Specifying the random seed in splitting the dataset.
+        Default: 717
+    train_size : float
+        training set size of the dataset.
+        Default: 0.7
+    val_size : float
+        validation set size of the dataset, and the
+        size of testing set is (1 - train_size - val_size)
+        Default: 0.1
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of label classes
+    graph : dgl.DGLGraph
+        Graph structure, etc.
+    seed : int
+        Random seed in splitting the dataset.
+    train_size : float
+        Training set size of the dataset.
+    val_size : float
+        Validation set size of the dataset
+
+    Examples
+    --------
+    >>> dataset = FraudDataset('yelp')
+    >>> graph = dataset[0]
+    >>> num_classes = dataset.num_classes
+    >>> feat = graph.ndata['feature']
+    >>> label = graph.ndata['label']
+    """
+    file_urls = {
+        "yelp": "dataset/FraudYelp.zip",
+        "amazon": "dataset/FraudAmazon.zip",
+    }
+    relations = {
+        "yelp": ["net_rsr", "net_rtr", "net_rur"],
+        "amazon": ["net_upu", "net_usu", "net_uvu"],
+    }
+    file_names = {"yelp": "YelpChi.mat", "amazon": "Amazon.mat"}
+    node_name = {"yelp": "review", "amazon": "user"}
+
+    def __init__(
+        self,
+        name,
+        raw_dir=None,
+        random_seed=717,
+        train_size=0.7,
+        val_size=0.1,
+        force_reload=False,
+        verbose=True,
+        transform=None,
+    ):
+        assert name in ["yelp", "amazon"], "only supports 'yelp', or 'amazon'"
+        url = _get_dgl_url(self.file_urls[name])
+        self.seed = random_seed
+        self.train_size = train_size
+        self.val_size = val_size
+        super(FraudDataset, self).__init__(
+            name=name,
+            url=url,
+            raw_dir=raw_dir,
+            hash_key=(random_seed, train_size, val_size),
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        """process raw data to graph, labels, splitting masks"""
+        file_path = os.path.join(self.raw_path, self.file_names[self.name])
+
+        data = io.loadmat(file_path)
+        node_features = data["features"].todense()
+        # remove additional dimension of length 1 in raw .mat file
+        node_labels = data["label"].squeeze()
+
+        graph_data = {}
+        for relation in self.relations[self.name]:
+            adj = data[relation].tocoo()
+            row, col = adj.row, adj.col
+            graph_data[
+                (self.node_name[self.name], relation, self.node_name[self.name])
+            ] = (row, col)
+        g = heterograph(graph_data)
+
+        g.ndata["feature"] = F.tensor(
+            node_features, dtype=F.data_type_dict["float32"]
+        )
+        g.ndata["label"] = F.tensor(
+            node_labels, dtype=F.data_type_dict["int64"]
+        )
+        self.graph = g
+
+        self._random_split(
+            g.ndata["feature"], self.seed, self.train_size, self.val_size
+        )
+
+    def __getitem__(self, idx):
+        r"""Get graph object
+
+        Parameters
+        ----------
+        idx : int
+            Item index
+
+        Returns
+        -------
+        :class:`dgl.DGLGraph`
+            graph structure, node features, node labels and masks
+
+            - ``ndata['feature']``: node features
+            - ``ndata['label']``: node labels
+            - ``ndata['train_mask']``: mask of training set
+            - ``ndata['val_mask']``: mask of validation set
+            - ``ndata['test_mask']``: mask of testing set
+        """
+        assert idx == 0, "This dataset has only one graph"
+        if self._transform is None:
+            return self.graph
+        else:
+            return self._transform(self.graph)
+
+    def __len__(self):
+        """number of data examples"""
+        return len(self.graph)
+
+    @property
+    def num_classes(self):
+        """Number of classes.
+
+        Return
+        -------
+        int
+        """
+        return 2
+
+    def save(self):
+        """save processed data to directory `self.save_path`"""
+        graph_path = os.path.join(
+            self.save_path, self.name + "_dgl_graph_{}.bin".format(self.hash)
+        )
+        save_graphs(str(graph_path), self.graph)
+
+    def load(self):
+        """load processed data from directory `self.save_path`"""
+        graph_path = os.path.join(
+            self.save_path, self.name + "_dgl_graph_{}.bin".format(self.hash)
+        )
+        graph_list, _ = load_graphs(str(graph_path))
+        g = graph_list[0]
+        self.graph = g
+
+    def has_cache(self):
+        """check whether there are processed data in `self.save_path`"""
+        graph_path = os.path.join(
+            self.save_path, self.name + "_dgl_graph_{}.bin".format(self.hash)
+        )
+        return os.path.exists(graph_path)
+
+    def _random_split(self, x, seed=717, train_size=0.7, val_size=0.1):
+        """split the dataset into training set, validation set and testing set"""
+
+        assert 0 <= train_size + val_size <= 1, (
+            "The sum of valid training set size and validation set size "
+            "must between 0 and 1 (inclusive)."
+        )
+
+        N = x.shape[0]
+        index = np.arange(N)
+        if self.name == "amazon":
+            # 0-3304 are unlabeled nodes
+            index = np.arange(3305, N)
+
+        index = np.random.RandomState(seed).permutation(index)
+        train_idx = index[: int(train_size * len(index))]
+        val_idx = index[len(index) - int(val_size * len(index)) :]
+        test_idx = index[
+            int(train_size * len(index)) : len(index)
+            - int(val_size * len(index))
+        ]
+        train_mask = np.zeros(N, dtype=np.bool_)
+        val_mask = np.zeros(N, dtype=np.bool_)
+        test_mask = np.zeros(N, dtype=np.bool_)
+        train_mask[train_idx] = True
+        val_mask[val_idx] = True
+        test_mask[test_idx] = True
+        self.graph.ndata["train_mask"] = F.tensor(train_mask)
+        self.graph.ndata["val_mask"] = F.tensor(val_mask)
+        self.graph.ndata["test_mask"] = F.tensor(test_mask)
+
+
+class FraudYelpDataset(FraudDataset):
+    r"""Fraud Yelp Dataset
+
+    The Yelp dataset includes hotel and restaurant reviews filtered (spam) and recommended
+    (legitimate) by Yelp. A spam review detection task can be conducted, which is a binary
+    classification task. 32 handcrafted features from <http://dx.doi.org/10.1145/2783258.2783370>
+    are taken as the raw node features. Reviews are nodes in the graph, and three relations are:
+
+        1. R-U-R: it connects reviews posted by the same user
+        2. R-S-R: it connects reviews under the same product with the same star rating (1-5 stars)
+        3. R-T-R: it connects two reviews under the same product posted in the same month.
+
+    Statistics:
+
+    - Nodes: 45,954
+    - Edges:
+
+        - R-U-R: 98,630
+        - R-T-R: 1,147,232
+        - R-S-R: 6,805,486
+
+    - Classes:
+
+        - Positive (spam): 6,677
+        - Negative (legitimate): 39,277
+
+    - Positive-Negative ratio: 1 : 5.9
+    - Node feature size: 32
+
+    Parameters
+    ----------
+    raw_dir : str
+        Specifying the directory that will store the
+        downloaded data or the directory that
+        already stores the input data.
+        Default: ~/.dgl/
+    random_seed : int
+        Specifying the random seed in splitting the dataset.
+        Default: 717
+    train_size : float
+        training set size of the dataset.
+        Default: 0.7
+    val_size : float
+        validation set size of the dataset, and the
+        size of testing set is (1 - train_size - val_size)
+        Default: 0.1
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Examples
+    --------
+    >>> dataset = FraudYelpDataset()
+    >>> graph = dataset[0]
+    >>> num_classes = dataset.num_classes
+    >>> feat = graph.ndata['feature']
+    >>> label = graph.ndata['label']
+    """
+
+    def __init__(
+        self,
+        raw_dir=None,
+        random_seed=717,
+        train_size=0.7,
+        val_size=0.1,
+        force_reload=False,
+        verbose=True,
+        transform=None,
+    ):
+        super(FraudYelpDataset, self).__init__(
+            name="yelp",
+            raw_dir=raw_dir,
+            random_seed=random_seed,
+            train_size=train_size,
+            val_size=val_size,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+
+class FraudAmazonDataset(FraudDataset):
+    r"""Fraud Amazon Dataset
+
+    The Amazon dataset includes product reviews under the Musical Instruments category.
+    Users with more than 80% helpful votes are labelled as benign entities and users with
+    less than 20% helpful votes are labelled as fraudulent entities. A fraudulent user
+    detection task can be conducted on the Amazon dataset, which is a binary classification
+    task. 25 handcrafted features from <https://arxiv.org/pdf/2005.10150.pdf> are taken as
+    the raw node features .
+
+    Users are nodes in the graph, and three relations are:
+    1. U-P-U : it connects users reviewing at least one same product
+    2. U-S-U : it connects users having at least one same star rating within one week
+    3. U-V-U : it connects users with top 5% mutual review text similarities (measured by
+    TF-IDF) among all users.
+
+    Statistics:
+
+    - Nodes: 11,944
+    - Edges:
+
+        - U-P-U: 351,216
+        - U-S-U: 7,132,958
+        - U-V-U: 2,073,474
+
+    - Classes:
+
+        - Positive (fraudulent): 821
+        - Negative (benign): 7,818
+        - Unlabeled: 3,305
+
+    - Positive-Negative ratio: 1 : 10.5
+    - Node feature size: 25
+
+    Parameters
+    ----------
+    raw_dir : str
+        Specifying the directory that will store the
+        downloaded data or the directory that
+        already stores the input data.
+        Default: ~/.dgl/
+    random_seed : int
+        Specifying the random seed in splitting the dataset.
+        Default: 717
+    train_size : float
+        training set size of the dataset.
+        Default: 0.7
+    val_size : float
+        validation set size of the dataset, and the
+        size of testing set is (1 - train_size - val_size)
+        Default: 0.1
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Examples
+    --------
+    >>> dataset = FraudAmazonDataset()
+    >>> graph = dataset[0]
+    >>> num_classes = dataset.num_classes
+    >>> feat = graph.ndata['feature']
+    >>> label = graph.ndata['label']
+    """
+
+    def __init__(
+        self,
+        raw_dir=None,
+        random_seed=717,
+        train_size=0.7,
+        val_size=0.1,
+        force_reload=False,
+        verbose=True,
+        transform=None,
+    ):
+        super(FraudAmazonDataset, self).__init__(
+            name="amazon",
+            raw_dir=raw_dir,
+            random_seed=random_seed,
+            train_size=train_size,
+            val_size=val_size,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )

python/dgl/data/gdelt.py

@@ -0,0 +1,203 @@
+""" GDELT dataset for temporal graph """
+import os
+
+import numpy as np
+
+from .. import backend as F
+from ..convert import graph as dgl_graph
+from .dgl_dataset import DGLBuiltinDataset
+from .utils import _get_dgl_url, load_info, loadtxt, save_info
+
+
+class GDELTDataset(DGLBuiltinDataset):
+    r"""GDELT dataset for event-based temporal graph
+
+    The Global Database of Events, Language, and Tone (GDELT) dataset.
+    This contains events happend all over the world (ie every protest held
+    anywhere in Russia on a given day is collapsed to a single entry).
+    This Dataset consists ofevents collected from 1/1/2018 to 1/31/2018
+    (15 minutes time granularity).
+
+    Reference:
+
+        - `Recurrent Event Network for Reasoning over Temporal Knowledge Graphs <https://arxiv.org/abs/1904.05530>`_
+        - `The Global Database of Events, Language, and Tone (GDELT) <https://dataverse.harvard.edu/dataset.xhtml?persistentId=doi:10.7910/DVN/28075>`_
+
+    Statistics:
+
+    - Train examples: 2,304
+    - Valid examples: 288
+    - Test examples: 384
+
+    Parameters
+    ----------
+    mode : str
+        Must be one of ('train', 'valid', 'test'). Default: 'train'
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    start_time : int
+        Start time of the temporal graph
+    end_time : int
+        End time of the temporal graph
+    is_temporal : bool
+        Does the dataset contain temporal graphs
+
+    Examples
+    ----------
+    >>> # get train, valid, test dataset
+    >>> train_data = GDELTDataset()
+    >>> valid_data = GDELTDataset(mode='valid')
+    >>> test_data = GDELTDataset(mode='test')
+    >>>
+    >>> # length of train set
+    >>> train_size = len(train_data)
+    >>>
+    >>> for g in train_data:
+    ....    e_feat = g.edata['rel_type']
+    ....    # your code here
+    ....
+    >>>
+    """
+
+    def __init__(
+        self,
+        mode="train",
+        raw_dir=None,
+        force_reload=False,
+        verbose=False,
+        transform=None,
+    ):
+        mode = mode.lower()
+        assert mode in ["train", "valid", "test"], "Mode not valid."
+        self.mode = mode
+        self.num_nodes = 23033
+        _url = _get_dgl_url("dataset/gdelt.zip")
+        super(GDELTDataset, self).__init__(
+            name="GDELT",
+            url=_url,
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        file_path = os.path.join(self.raw_path, self.mode + ".txt")
+        self.data = loadtxt(file_path, delimiter="\t").astype(np.int64)
+
+        # The source code is not released, but the paper indicates there're
+        # totally 137 samples. The cutoff below has exactly 137 samples.
+        self.time_index = np.floor(self.data[:, 3] / 15).astype(np.int64)
+        self._start_time = self.time_index.min()
+        self._end_time = self.time_index.max()
+
+    @property
+    def info_path(self):
+        return os.path.join(self.save_path, self.mode + "_info.pkl")
+
+    def has_cache(self):
+        return os.path.exists(self.info_path)
+
+    def save(self):
+        save_info(
+            self.info_path,
+            {
+                "data": self.data,
+                "time_index": self.time_index,
+                "start_time": self.start_time,
+                "end_time": self.end_time,
+            },
+        )
+
+    def load(self):
+        info = load_info(self.info_path)
+        self.data, self.time_index, self._start_time, self._end_time = (
+            info["data"],
+            info["time_index"],
+            info["start_time"],
+            info["end_time"],
+        )
+
+    @property
+    def start_time(self):
+        r"""Start time of events in the temporal graph
+
+        Returns
+        -------
+        int
+        """
+        return self._start_time
+
+    @property
+    def end_time(self):
+        r"""End time of events in the temporal graph
+
+        Returns
+        -------
+        int
+        """
+        return self._end_time
+
+    def __getitem__(self, t):
+        r"""Get graph by with events before time `t + self.start_time`
+
+        Parameters
+        ----------
+        t : int
+            Time, its value must be in range [0, `self.end_time` - `self.start_time`]
+
+        Returns
+        -------
+        :class:`dgl.DGLGraph`
+
+            The graph contains:
+
+            - ``edata['rel_type']``: edge type
+        """
+        if t >= len(self) or t < 0:
+            raise IndexError("Index out of range")
+        i = t + self.start_time
+        row_mask = self.time_index <= i
+        edges = self.data[row_mask][:, [0, 2]]
+        rate = self.data[row_mask][:, 1]
+        g = dgl_graph((edges[:, 0], edges[:, 1]))
+        g.edata["rel_type"] = F.tensor(
+            rate.reshape(-1, 1), dtype=F.data_type_dict["int64"]
+        )
+        if self._transform is not None:
+            g = self._transform(g)
+        return g
+
+    def __len__(self):
+        r"""Number of graphs in the dataset.
+
+        Return
+        -------
+        int
+        """
+        return self._end_time - self._start_time + 1
+
+    @property
+    def is_temporal(self):
+        r"""Does the dataset contain temporal graphs
+
+        Returns
+        -------
+        bool
+        """
+        return True
+
+
+GDELT = GDELTDataset

python/dgl/data/geom_gcn.py

@@ -0,0 +1,483 @@
+"""Datasets introduced in the Geom-GCN paper."""
+import os
+
+import numpy as np
+
+from ..convert import graph
+from .dgl_dataset import DGLBuiltinDataset
+from .utils import _get_dgl_url
+
+
+class GeomGCNDataset(DGLBuiltinDataset):
+    r"""Datasets introduced in
+    `Geom-GCN: Geometric Graph Convolutional Networks
+    <https://arxiv.org/abs/2002.05287>`__
+
+    Parameters
+    ----------
+    name : str
+        Name of the dataset.
+    raw_dir : str
+        Raw file directory to store the processed data.
+    force_reload : bool
+        Whether to re-download the data source.
+    verbose : bool
+        Whether to print progress information.
+    transform : callable
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+    """
+
+    def __init__(self, name, raw_dir, force_reload, verbose, transform):
+        url = _get_dgl_url(f"dataset/{name}.zip")
+        super(GeomGCNDataset, self).__init__(
+            name=name,
+            url=url,
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        """Load and process the data."""
+        try:
+            import torch
+        except ImportError:
+            raise ModuleNotFoundError(
+                "This dataset requires PyTorch to be the backend."
+            )
+
+        # Process node features and labels.
+        with open(f"{self.raw_path}/out1_node_feature_label.txt", "r") as f:
+            data = f.read().split("\n")[1:-1]
+        features = [
+            [float(v) for v in r.split("\t")[1].split(",")] for r in data
+        ]
+        features = torch.tensor(features, dtype=torch.float)
+        labels = [int(r.split("\t")[2]) for r in data]
+        self._num_classes = max(labels) + 1
+        labels = torch.tensor(labels, dtype=torch.long)
+
+        # Process graph structure.
+        with open(f"{self.raw_path}/out1_graph_edges.txt", "r") as f:
+            data = f.read().split("\n")[1:-1]
+            data = [[int(v) for v in r.split("\t")] for r in data]
+        dst, src = torch.tensor(data, dtype=torch.long).t().contiguous()
+
+        self._g = graph((src, dst), num_nodes=features.size(0))
+        self._g.ndata["feat"] = features
+        self._g.ndata["label"] = labels
+
+        # Process 10 train/val/test node splits.
+        train_masks, val_masks, test_masks = [], [], []
+        for i in range(10):
+            filepath = f"{self.raw_path}/{self.name}_split_0.6_0.2_{i}.npz"
+            f = np.load(filepath)
+            train_masks += [torch.from_numpy(f["train_mask"])]
+            val_masks += [torch.from_numpy(f["val_mask"])]
+            test_masks += [torch.from_numpy(f["test_mask"])]
+        self._g.ndata["train_mask"] = torch.stack(train_masks, dim=1).bool()
+        self._g.ndata["val_mask"] = torch.stack(val_masks, dim=1).bool()
+        self._g.ndata["test_mask"] = torch.stack(test_masks, dim=1).bool()
+
+    def has_cache(self):
+        return os.path.exists(self.raw_path)
+
+    def load(self):
+        self.process()
+
+    def __getitem__(self, idx):
+        assert idx == 0, "This dataset has only one graph."
+        if self._transform is None:
+            return self._g
+        else:
+            return self._transform(self._g)
+
+    def __len__(self):
+        return 1
+
+    @property
+    def num_classes(self):
+        return self._num_classes
+
+
+class ChameleonDataset(GeomGCNDataset):
+    r"""Wikipedia page-page network on chameleons from `Multi-scale Attributed
+    Node Embedding <https://arxiv.org/abs/1909.13021>`__ and later modified by
+    `Geom-GCN: Geometric Graph Convolutional Networks
+    <https://arxiv.org/abs/2002.05287>`__
+
+    Nodes represent articles from the English Wikipedia, edges reflect mutual
+    links between them. Node features indicate the presence of particular nouns
+    in the articles. The nodes were classified into 5 classes in terms of their
+    average monthly traffic.
+
+    Statistics:
+
+    - Nodes: 2277
+    - Edges: 36101
+    - Number of Classes: 5
+    - 10 train/val/test splits
+
+        - Train: 1092
+        - Val: 729
+        - Test: 456
+
+    Parameters
+    ----------
+    raw_dir : str, optional
+        Raw file directory to store the processed data. Default: ~/.dgl/
+    force_reload : bool, optional
+        Whether to re-download the data source. Default: False
+    verbose : bool, optional
+        Whether to print progress information. Default: True
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access. Default: None
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of node classes
+
+    Notes
+    -----
+    The graph does not come with edges for both directions.
+
+    Examples
+    --------
+
+    >>> from dgl.data import ChameleonDataset
+    >>> dataset = ChameleonDataset()
+    >>> g = dataset[0]
+    >>> num_classes = dataset.num_classes
+
+    >>> # get node features
+    >>> feat = g.ndata["feat"]
+
+    >>> # get data split
+    >>> train_mask = g.ndata["train_mask"]
+    >>> val_mask = g.ndata["val_mask"]
+    >>> test_mask = g.ndata["test_mask"]
+
+    >>> # get labels
+    >>> label = g.ndata['label']
+    """
+
+    def __init__(
+        self, raw_dir=None, force_reload=False, verbose=True, transform=None
+    ):
+        super(ChameleonDataset, self).__init__(
+            name="chameleon",
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+
+class SquirrelDataset(GeomGCNDataset):
+    r"""Wikipedia page-page network on squirrels from `Multi-scale Attributed
+    Node Embedding <https://arxiv.org/abs/1909.13021>`__ and later modified by
+    `Geom-GCN: Geometric Graph Convolutional Networks
+    <https://arxiv.org/abs/2002.05287>`__
+
+    Nodes represent articles from the English Wikipedia, edges reflect mutual
+    links between them. Node features indicate the presence of particular nouns
+    in the articles. The nodes were classified into 5 classes in terms of their
+    average monthly traffic.
+
+    Statistics:
+
+    - Nodes: 5201
+    - Edges: 217073
+    - Number of Classes: 5
+    - 10 train/val/test splits
+
+        - Train: 2496
+        - Val: 1664
+        - Test: 1041
+
+    Parameters
+    ----------
+    raw_dir : str, optional
+        Raw file directory to store the processed data. Default: ~/.dgl/
+    force_reload : bool, optional
+        Whether to re-download the data source. Default: False
+    verbose : bool, optional
+        Whether to print progress information. Default: True
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access. Default: None
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of node classes
+
+    Notes
+    -----
+    The graph does not come with edges for both directions.
+
+    Examples
+    --------
+
+    >>> from dgl.data import SquirrelDataset
+    >>> dataset = SquirrelDataset()
+    >>> g = dataset[0]
+    >>> num_classes = dataset.num_classes
+
+    >>> # get node features
+    >>> feat = g.ndata["feat"]
+
+    >>> # get data split
+    >>> train_mask = g.ndata["train_mask"]
+    >>> val_mask = g.ndata["val_mask"]
+    >>> test_mask = g.ndata["test_mask"]
+
+    >>> # get labels
+    >>> label = g.ndata['label']
+    """
+
+    def __init__(
+        self, raw_dir=None, force_reload=False, verbose=True, transform=None
+    ):
+        super(SquirrelDataset, self).__init__(
+            name="squirrel",
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+
+class CornellDataset(GeomGCNDataset):
+    r"""Cornell subset of
+    `WebKB <http://www.cs.cmu.edu/afs/cs.cmu.edu/project/theo-11/www/wwkb/>`__,
+    later modified by `Geom-GCN: Geometric Graph Convolutional Networks
+    <https://arxiv.org/abs/2002.05287>`__
+
+    Nodes represent web pages. Edges represent hyperlinks between them. Node
+    features are the bag-of-words representation of web pages. The web pages
+    are manually classified into the five categories, student, project, course,
+    staff, and faculty.
+
+    Statistics:
+
+    - Nodes: 183
+    - Edges: 298
+    - Number of Classes: 5
+    - 10 train/val/test splits
+
+        - Train: 87
+        - Val: 59
+        - Test: 37
+
+    Parameters
+    ----------
+    raw_dir : str, optional
+        Raw file directory to store the processed data. Default: ~/.dgl/
+    force_reload : bool, optional
+        Whether to re-download the data source. Default: False
+    verbose : bool, optional
+        Whether to print progress information. Default: True
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access. Default: None
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of node classes
+
+    Notes
+    -----
+    The graph does not come with edges for both directions.
+
+    Examples
+    --------
+
+    >>> from dgl.data import CornellDataset
+    >>> dataset = CornellDataset()
+    >>> g = dataset[0]
+    >>> num_classes = dataset.num_classes
+
+    >>> # get node features
+    >>> feat = g.ndata["feat"]
+
+    >>> # get data split
+    >>> train_mask = g.ndata["train_mask"]
+    >>> val_mask = g.ndata["val_mask"]
+    >>> test_mask = g.ndata["test_mask"]
+
+    >>> # get labels
+    >>> label = g.ndata['label']
+    """
+
+    def __init__(
+        self, raw_dir=None, force_reload=False, verbose=True, transform=None
+    ):
+        super(CornellDataset, self).__init__(
+            name="cornell",
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+
+class TexasDataset(GeomGCNDataset):
+    r"""Texas subset of
+    `WebKB <http://www.cs.cmu.edu/afs/cs.cmu.edu/project/theo-11/www/wwkb/>`__,
+    later modified by `Geom-GCN: Geometric Graph Convolutional Networks
+    <https://arxiv.org/abs/2002.05287>`__
+
+    Nodes represent web pages. Edges represent hyperlinks between them. Node
+    features are the bag-of-words representation of web pages. The web pages
+    are manually classified into the five categories, student, project, course,
+    staff, and faculty.
+
+    Statistics:
+
+    - Nodes: 183
+    - Edges: 325
+    - Number of Classes: 5
+    - 10 train/val/test splits
+
+        - Train: 87
+        - Val: 59
+        - Test: 37
+
+    Parameters
+    ----------
+    raw_dir : str, optional
+        Raw file directory to store the processed data. Default: ~/.dgl/
+    force_reload : bool, optional
+        Whether to re-download the data source. Default: False
+    verbose : bool, optional
+        Whether to print progress information. Default: True
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access. Default: None
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of node classes
+
+    Notes
+    -----
+    The graph does not come with edges for both directions.
+
+    Examples
+    --------
+
+    >>> from dgl.data import TexasDataset
+    >>> dataset = TexasDataset()
+    >>> g = dataset[0]
+    >>> num_classes = dataset.num_classes
+
+    >>> # get node features
+    >>> feat = g.ndata["feat"]
+
+    >>> # get data split
+    >>> train_mask = g.ndata["train_mask"]
+    >>> val_mask = g.ndata["val_mask"]
+    >>> test_mask = g.ndata["test_mask"]
+
+    >>> # get labels
+    >>> label = g.ndata['label']
+    """
+
+    def __init__(
+        self, raw_dir=None, force_reload=False, verbose=True, transform=None
+    ):
+        super(TexasDataset, self).__init__(
+            name="texas",
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+
+class WisconsinDataset(GeomGCNDataset):
+    r"""Wisconsin subset of
+    `WebKB <http://www.cs.cmu.edu/afs/cs.cmu.edu/project/theo-11/www/wwkb/>`__,
+    later modified by `Geom-GCN: Geometric Graph Convolutional Networks
+    <https://arxiv.org/abs/2002.05287>`__
+
+    Nodes represent web pages. Edges represent hyperlinks between them. Node
+    features are the bag-of-words representation of web pages. The web pages
+    are manually classified into the five categories, student, project, course,
+    staff, and faculty.
+
+    Statistics:
+
+    - Nodes: 251
+    - Edges: 515
+    - Number of Classes: 5
+    - 10 train/val/test splits
+
+        - Train: 120
+        - Val: 80
+        - Test: 51
+
+    Parameters
+    ----------
+    raw_dir : str, optional
+        Raw file directory to store the processed data. Default: ~/.dgl/
+    force_reload : bool, optional
+        Whether to re-download the data source. Default: False
+    verbose : bool, optional
+        Whether to print progress information. Default: True
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access. Default: None
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of node classes
+
+    Notes
+    -----
+    The graph does not come with edges for both directions.
+
+    Examples
+    --------
+
+    >>> from dgl.data import WisconsinDataset
+    >>> dataset = WisconsinDataset()
+    >>> g = dataset[0]
+    >>> num_classes = dataset.num_classes
+
+    >>> # get node features
+    >>> feat = g.ndata["feat"]
+
+    >>> # get data split
+    >>> train_mask = g.ndata["train_mask"]
+    >>> val_mask = g.ndata["val_mask"]
+    >>> test_mask = g.ndata["test_mask"]
+
+    >>> # get labels
+    >>> label = g.ndata['label']
+    """
+
+    def __init__(
+        self, raw_dir=None, force_reload=False, verbose=True, transform=None
+    ):
+        super(WisconsinDataset, self).__init__(
+            name="wisconsin",
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )

python/dgl/data/gindt.py

@@ -0,0 +1,420 @@
+"""Datasets used in How Powerful Are Graph Neural Networks?
+(chen jun)
+Datasets include:
+MUTAG, COLLAB, IMDBBINARY, IMDBMULTI, NCI1, PROTEINS, PTC, REDDITBINARY, REDDITMULTI5K
+https://github.com/weihua916/powerful-gnns/blob/master/dataset.zip
+"""
+
+import os
+
+import numpy as np
+
+from .. import backend as F
+from ..convert import graph as dgl_graph
+from ..utils import retry_method_with_fix
+from .dgl_dataset import DGLBuiltinDataset
+from .utils import (
+    download,
+    extract_archive,
+    load_graphs,
+    load_info,
+    loadtxt,
+    save_graphs,
+    save_info,
+)
+
+
+class GINDataset(DGLBuiltinDataset):
+    """Dataset Class for `How Powerful Are Graph Neural Networks? <https://arxiv.org/abs/1810.00826>`_.
+
+    This is adapted from `<https://github.com/weihua916/powerful-gnns/blob/master/dataset.zip>`_.
+
+    The class provides an interface for nine datasets used in the paper along with the paper-specific
+    settings. The datasets are ``'MUTAG'``, ``'COLLAB'``, ``'IMDBBINARY'``, ``'IMDBMULTI'``,
+    ``'NCI1'``, ``'PROTEINS'``, ``'PTC'``, ``'REDDITBINARY'``, ``'REDDITMULTI5K'``.
+
+    If ``degree_as_nlabel`` is set to ``False``, then ``ndata['label']`` stores the provided node label,
+    otherwise ``ndata['label']`` stores the node in-degrees.
+
+    For graphs that have node attributes, ``ndata['attr']`` stores the node attributes.
+    For graphs that have no attribute, ``ndata['attr']`` stores the corresponding one-hot encoding
+    of ``ndata['label']``.
+
+    Parameters
+    ---------
+    name: str
+        dataset name, one of
+        (``'MUTAG'``, ``'COLLAB'``, \
+        ``'IMDBBINARY'``, ``'IMDBMULTI'``, \
+        ``'NCI1'``, ``'PROTEINS'``, ``'PTC'``, \
+        ``'REDDITBINARY'``, ``'REDDITMULTI5K'``)
+    self_loop: bool
+        add self to self edge if true
+    degree_as_nlabel: bool
+        take node degree as label and feature if true
+    transform: callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of classes for multiclass classification
+
+    Examples
+    --------
+    >>> data = GINDataset(name='MUTAG', self_loop=False)
+
+    The dataset instance is an iterable
+
+    >>> len(data)
+    188
+    >>> g, label = data[128]
+    >>> g
+    Graph(num_nodes=13, num_edges=26,
+          ndata_schemes={'label': Scheme(shape=(), dtype=torch.int64), 'attr': Scheme(shape=(7,), dtype=torch.float32)}
+          edata_schemes={})
+    >>> label
+    tensor(1)
+
+    Batch the graphs and labels for mini-batch training
+
+    >>> graphs, labels = zip(*[data[i] for i in range(16)])
+    >>> batched_graphs = dgl.batch(graphs)
+    >>> batched_labels = torch.tensor(labels)
+    >>> batched_graphs
+    Graph(num_nodes=330, num_edges=748,
+          ndata_schemes={'label': Scheme(shape=(), dtype=torch.int64), 'attr': Scheme(shape=(7,), dtype=torch.float32)}
+          edata_schemes={})
+    """
+
+    def __init__(
+        self,
+        name,
+        self_loop,
+        degree_as_nlabel=False,
+        raw_dir=None,
+        force_reload=False,
+        verbose=False,
+        transform=None,
+    ):
+        self._name = name  # MUTAG
+        gin_url = "https://raw.githubusercontent.com/weihua916/powerful-gnns/master/dataset.zip"
+        self.ds_name = "nig"
+
+        self.self_loop = self_loop
+        self.graphs = []
+        self.labels = []
+
+        # relabel
+        self.glabel_dict = {}
+        self.nlabel_dict = {}
+        self.elabel_dict = {}
+        self.ndegree_dict = {}
+
+        # global num
+        self.N = 0  # total graphs number
+        self.n = 0  # total nodes number
+        self.m = 0  # total edges number
+
+        # global num of classes
+        self.gclasses = 0
+        self.nclasses = 0
+        self.eclasses = 0
+        self.dim_nfeats = 0
+
+        # flags
+        self.degree_as_nlabel = degree_as_nlabel
+        self.nattrs_flag = False
+        self.nlabels_flag = False
+
+        super(GINDataset, self).__init__(
+            name=name,
+            url=gin_url,
+            hash_key=(name, self_loop, degree_as_nlabel),
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    @property
+    def raw_path(self):
+        return os.path.join(self.raw_dir, "GINDataset")
+
+    def download(self):
+        r"""Automatically download data and extract it."""
+        zip_file_path = os.path.join(self.raw_dir, "GINDataset.zip")
+        download(self.url, path=zip_file_path)
+        extract_archive(zip_file_path, self.raw_path)
+
+    def __len__(self):
+        """Return the number of graphs in the dataset."""
+        return len(self.graphs)
+
+    def __getitem__(self, idx):
+        """Get the idx-th sample.
+
+        Parameters
+        ---------
+        idx : int
+            The sample index.
+
+        Returns
+        -------
+        (:class:`dgl.Graph`, Tensor)
+            The graph and its label.
+        """
+        if self._transform is None:
+            g = self.graphs[idx]
+        else:
+            g = self._transform(self.graphs[idx])
+        return g, self.labels[idx]
+
+    def _file_path(self):
+        return os.path.join(
+            self.raw_dir,
+            "GINDataset",
+            "dataset",
+            self.name,
+            "{}.txt".format(self.name),
+        )
+
+    def process(self):
+        """Loads input dataset from dataset/NAME/NAME.txt file"""
+        if self.verbose:
+            print("loading data...")
+        self.file = self._file_path()
+        with open(self.file, "r") as f:
+            # line_1 == N, total number of graphs
+            self.N = int(f.readline().strip())
+
+            for i in range(self.N):
+                if (i + 1) % 10 == 0 and self.verbose is True:
+                    print("processing graph {}...".format(i + 1))
+
+                grow = f.readline().strip().split()
+                # line_2 == [n_nodes, l] is equal to
+                # [node number of a graph, class label of a graph]
+                n_nodes, glabel = [int(w) for w in grow]
+
+                # relabel graphs
+                if glabel not in self.glabel_dict:
+                    mapped = len(self.glabel_dict)
+                    self.glabel_dict[glabel] = mapped
+
+                self.labels.append(self.glabel_dict[glabel])
+
+                g = dgl_graph(([], []))
+                g.add_nodes(n_nodes)
+
+                nlabels = []  # node labels
+                nattrs = []  # node attributes if it has
+                m_edges = 0
+
+                for j in range(n_nodes):
+                    nrow = f.readline().strip().split()
+
+                    # handle edges and attributes(if has)
+                    tmp = int(nrow[1]) + 2  # tmp == 2 + #edges
+                    if tmp == len(nrow):
+                        # no node attributes
+                        nrow = [int(w) for w in nrow]
+                    elif tmp > len(nrow):
+                        nrow = [int(w) for w in nrow[:tmp]]
+                        nattr = [float(w) for w in nrow[tmp:]]
+                        nattrs.append(nattr)
+                    else:
+                        raise Exception("edge number is incorrect!")
+
+                    # relabel nodes if it has labels
+                    # if it doesn't have node labels, then every nrow[0]==0
+                    if not nrow[0] in self.nlabel_dict:
+                        mapped = len(self.nlabel_dict)
+                        self.nlabel_dict[nrow[0]] = mapped
+
+                    nlabels.append(self.nlabel_dict[nrow[0]])
+
+                    m_edges += nrow[1]
+                    g.add_edges(j, nrow[2:])
+
+                    # add self loop
+                    if self.self_loop:
+                        m_edges += 1
+                        g.add_edges(j, j)
+
+                    if (j + 1) % 10 == 0 and self.verbose is True:
+                        print(
+                            "processing node {} of graph {}...".format(
+                                j + 1, i + 1
+                            )
+                        )
+                        print("this node has {} edgs.".format(nrow[1]))
+
+                if nattrs != []:
+                    nattrs = np.stack(nattrs)
+                    g.ndata["attr"] = F.tensor(nattrs, F.float32)
+                    self.nattrs_flag = True
+
+                g.ndata["label"] = F.tensor(nlabels)
+                if len(self.nlabel_dict) > 1:
+                    self.nlabels_flag = True
+
+                assert g.num_nodes() == n_nodes
+
+                # update statistics of graphs
+                self.n += n_nodes
+                self.m += m_edges
+
+                self.graphs.append(g)
+
+        self.labels = F.tensor(self.labels)
+        # if no attr
+        if not self.nattrs_flag:
+            if self.verbose:
+                print("there are no node features in this dataset!")
+            # generate node attr by node degree
+            if self.degree_as_nlabel:
+                if self.verbose:
+                    print("generate node features by node degree...")
+                for g in self.graphs:
+                    # actually this label shouldn't be updated
+                    # in case users want to keep it
+                    # but usually no features means no labels, fine.
+                    g.ndata["label"] = g.in_degrees()
+                    # extracting unique node labels
+
+            # in case the labels/degrees are not continuous number
+            nlabel_set = set([])
+            for g in self.graphs:
+                nlabel_set = nlabel_set.union(
+                    set([F.as_scalar(nl) for nl in g.ndata["label"]])
+                )
+            nlabel_set = list(nlabel_set)
+            is_label_valid = all(
+                [label in self.nlabel_dict for label in nlabel_set]
+            )
+            if (
+                is_label_valid
+                and len(nlabel_set) == np.max(nlabel_set) + 1
+                and np.min(nlabel_set) == 0
+            ):
+                # Note this is different from the author's implementation. In weihua916's implementation,
+                # the labels are relabeled anyway. But here we didn't relabel it if the labels are contiguous
+                # to make it consistent with the original dataset
+                label2idx = self.nlabel_dict
+            else:
+                label2idx = {nlabel_set[i]: i for i in range(len(nlabel_set))}
+            # generate node attr by node label
+            for g in self.graphs:
+                attr = np.zeros((g.num_nodes(), len(label2idx)))
+                attr[
+                    range(g.num_nodes()),
+                    [
+                        label2idx[nl]
+                        for nl in F.asnumpy(g.ndata["label"]).tolist()
+                    ],
+                ] = 1
+                g.ndata["attr"] = F.tensor(attr, F.float32)
+
+        # after load, get the #classes and #dim
+        self.gclasses = len(self.glabel_dict)
+        self.nclasses = len(self.nlabel_dict)
+        self.eclasses = len(self.elabel_dict)
+        self.dim_nfeats = len(self.graphs[0].ndata["attr"][0])
+
+        if self.verbose:
+            print("Done.")
+            print(
+                """
+                -------- Data Statistics --------'
+                #Graphs: %d
+                #Graph Classes: %d
+                #Nodes: %d
+                #Node Classes: %d
+                #Node Features Dim: %d
+                #Edges: %d
+                #Edge Classes: %d
+                Avg. of #Nodes: %.2f
+                Avg. of #Edges: %.2f
+                Graph Relabeled: %s
+                Node Relabeled: %s
+                Degree Relabeled(If degree_as_nlabel=True): %s \n """
+                % (
+                    self.N,
+                    self.gclasses,
+                    self.n,
+                    self.nclasses,
+                    self.dim_nfeats,
+                    self.m,
+                    self.eclasses,
+                    self.n / self.N,
+                    self.m / self.N,
+                    self.glabel_dict,
+                    self.nlabel_dict,
+                    self.ndegree_dict,
+                )
+            )
+
+    def save(self):
+        label_dict = {"labels": self.labels}
+        info_dict = {
+            "N": self.N,
+            "n": self.n,
+            "m": self.m,
+            "self_loop": self.self_loop,
+            "gclasses": self.gclasses,
+            "nclasses": self.nclasses,
+            "eclasses": self.eclasses,
+            "dim_nfeats": self.dim_nfeats,
+            "degree_as_nlabel": self.degree_as_nlabel,
+            "glabel_dict": self.glabel_dict,
+            "nlabel_dict": self.nlabel_dict,
+            "elabel_dict": self.elabel_dict,
+            "ndegree_dict": self.ndegree_dict,
+        }
+        save_graphs(str(self.graph_path), self.graphs, label_dict)
+        save_info(str(self.info_path), info_dict)
+
+    def load(self):
+        graphs, label_dict = load_graphs(str(self.graph_path))
+        info_dict = load_info(str(self.info_path))
+
+        self.graphs = graphs
+        self.labels = label_dict["labels"]
+
+        self.N = info_dict["N"]
+        self.n = info_dict["n"]
+        self.m = info_dict["m"]
+        self.self_loop = info_dict["self_loop"]
+        self.gclasses = info_dict["gclasses"]
+        self.nclasses = info_dict["nclasses"]
+        self.eclasses = info_dict["eclasses"]
+        self.dim_nfeats = info_dict["dim_nfeats"]
+        self.glabel_dict = info_dict["glabel_dict"]
+        self.nlabel_dict = info_dict["nlabel_dict"]
+        self.elabel_dict = info_dict["elabel_dict"]
+        self.ndegree_dict = info_dict["ndegree_dict"]
+        self.degree_as_nlabel = info_dict["degree_as_nlabel"]
+
+    @property
+    def graph_path(self):
+        return os.path.join(
+            self.save_path, "gin_{}_{}.bin".format(self.name, self.hash)
+        )
+
+    @property
+    def info_path(self):
+        return os.path.join(
+            self.save_path, "gin_{}_{}.pkl".format(self.name, self.hash)
+        )
+
+    def has_cache(self):
+        if os.path.exists(self.graph_path) and os.path.exists(self.info_path):
+            return True
+        return False
+
+    @property
+    def num_classes(self):
+        return self.gclasses

python/dgl/data/gnn_benchmark.py

@@ -0,0 +1,544 @@
+"""GNN Benchmark datasets for node classification."""
+import os
+
+import numpy as np
+import scipy.sparse as sp
+
+from .. import backend as F, transforms
+from ..convert import graph as dgl_graph
+
+from .dgl_dataset import DGLBuiltinDataset
+from .utils import (
+    _get_dgl_url,
+    deprecate_class,
+    deprecate_property,
+    load_graphs,
+    save_graphs,
+)
+
+__all__ = [
+    "AmazonCoBuyComputerDataset",
+    "AmazonCoBuyPhotoDataset",
+    "CoauthorPhysicsDataset",
+    "CoauthorCSDataset",
+    "CoraFullDataset",
+    "AmazonCoBuy",
+    "Coauthor",
+    "CoraFull",
+]
+
+
+def eliminate_self_loops(A):
+    """Remove self-loops from the adjacency matrix."""
+    A = A.tolil()
+    A.setdiag(0)
+    A = A.tocsr()
+    A.eliminate_zeros()
+    return A
+
+
+class GNNBenchmarkDataset(DGLBuiltinDataset):
+    r"""Base Class for GNN Benchmark dataset
+
+    Reference: https://github.com/shchur/gnn-benchmark#datasets
+    """
+
+    def __init__(
+        self,
+        name,
+        raw_dir=None,
+        force_reload=False,
+        verbose=False,
+        transform=None,
+    ):
+        _url = _get_dgl_url("dataset/" + name + ".zip")
+        super(GNNBenchmarkDataset, self).__init__(
+            name=name,
+            url=_url,
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        npz_path = os.path.join(self.raw_path, self.name + ".npz")
+        g = self._load_npz(npz_path)
+        g = transforms.reorder_graph(
+            g,
+            node_permute_algo="rcmk",
+            edge_permute_algo="dst",
+            store_ids=False,
+        )
+        self._graph = g
+        self._data = [g]
+        self._print_info()
+
+    def has_cache(self):
+        graph_path = os.path.join(self.save_path, "dgl_graph_v1.bin")
+        if os.path.exists(graph_path):
+            return True
+        return False
+
+    def save(self):
+        graph_path = os.path.join(self.save_path, "dgl_graph_v1.bin")
+        save_graphs(graph_path, self._graph)
+
+    def load(self):
+        graph_path = os.path.join(self.save_path, "dgl_graph_v1.bin")
+        graphs, _ = load_graphs(graph_path)
+        self._graph = graphs[0]
+        self._data = [graphs[0]]
+        self._print_info()
+
+    def _print_info(self):
+        if self.verbose:
+            print("  NumNodes: {}".format(self._graph.num_nodes()))
+            print("  NumEdges: {}".format(self._graph.num_edges()))
+            print("  NumFeats: {}".format(self._graph.ndata["feat"].shape[-1]))
+            print("  NumbClasses: {}".format(self.num_classes))
+
+    def _load_npz(self, file_name):
+        with np.load(file_name, allow_pickle=True) as loader:
+            loader = dict(loader)
+            num_nodes = loader["adj_shape"][0]
+            adj_matrix = sp.csr_matrix(
+                (
+                    loader["adj_data"],
+                    loader["adj_indices"],
+                    loader["adj_indptr"],
+                ),
+                shape=loader["adj_shape"],
+            ).tocoo()
+
+            if "attr_data" in loader:
+                # Attributes are stored as a sparse CSR matrix
+                attr_matrix = sp.csr_matrix(
+                    (
+                        loader["attr_data"],
+                        loader["attr_indices"],
+                        loader["attr_indptr"],
+                    ),
+                    shape=loader["attr_shape"],
+                ).todense()
+            elif "attr_matrix" in loader:
+                # Attributes are stored as a (dense) np.ndarray
+                attr_matrix = loader["attr_matrix"]
+            else:
+                attr_matrix = None
+
+            if "labels_data" in loader:
+                # Labels are stored as a CSR matrix
+                labels = sp.csr_matrix(
+                    (
+                        loader["labels_data"],
+                        loader["labels_indices"],
+                        loader["labels_indptr"],
+                    ),
+                    shape=loader["labels_shape"],
+                ).todense()
+            elif "labels" in loader:
+                # Labels are stored as a numpy array
+                labels = loader["labels"]
+            else:
+                labels = None
+        g = dgl_graph((adj_matrix.row, adj_matrix.col))
+        g = transforms.to_bidirected(g)
+        g.ndata["feat"] = F.tensor(attr_matrix, F.data_type_dict["float32"])
+        g.ndata["label"] = F.tensor(labels, F.data_type_dict["int64"])
+        return g
+
+    @property
+    def num_classes(self):
+        """Number of classes."""
+        raise NotImplementedError
+
+    def __getitem__(self, idx):
+        r"""Get graph by index
+
+        Parameters
+        ----------
+        idx : int
+            Item index
+
+        Returns
+        -------
+        :class:`dgl.DGLGraph`
+
+            The graph contains:
+
+            - ``ndata['feat']``: node features
+            - ``ndata['label']``: node labels
+        """
+        assert idx == 0, "This dataset has only one graph"
+        if self._transform is None:
+            return self._graph
+        else:
+            return self._transform(self._graph)
+
+    def __len__(self):
+        r"""Number of graphs in the dataset"""
+        return 1
+
+
+class CoraFullDataset(GNNBenchmarkDataset):
+    r"""CORA-Full dataset for node classification task.
+
+    Extended Cora dataset. Nodes represent paper and edges represent citations.
+
+    Reference: `<https://github.com/shchur/gnn-benchmark#datasets>`_
+
+    Statistics:
+
+    - Nodes: 19,793
+    - Edges: 126,842 (note that the original dataset has 65,311 edges but DGL adds
+      the reverse edges and remove the duplicates, hence with a different number)
+    - Number of Classes: 70
+    - Node feature size: 8,710
+
+    Parameters
+    ----------
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of classes for each node.
+
+    Examples
+    --------
+    >>> data = CoraFullDataset()
+    >>> g = data[0]
+    >>> num_class = data.num_classes
+    >>> feat = g.ndata['feat']  # get node feature
+    >>> label = g.ndata['label']  # get node labels
+    """
+
+    def __init__(
+        self, raw_dir=None, force_reload=False, verbose=False, transform=None
+    ):
+        super(CoraFullDataset, self).__init__(
+            name="cora_full",
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    @property
+    def num_classes(self):
+        """Number of classes.
+
+        Return
+        -------
+        int
+        """
+        return 70
+
+
+class CoauthorCSDataset(GNNBenchmarkDataset):
+    r"""'Computer Science (CS)' part of the Coauthor dataset for node classification task.
+
+    Coauthor CS and Coauthor Physics are co-authorship graphs based on the Microsoft Academic Graph
+    from the KDD Cup 2016 challenge. Here, nodes are authors, that are connected by an edge if they
+    co-authored a paper; node features represent paper keywords for each author’s papers, and class
+    labels indicate most active fields of study for each author.
+
+    Reference: `<https://github.com/shchur/gnn-benchmark#datasets>`_
+
+    Statistics:
+
+    - Nodes: 18,333
+    - Edges: 163,788 (note that the original dataset has 81,894 edges but DGL adds
+      the reverse edges and remove the duplicates, hence with a different number)
+    - Number of classes: 15
+    - Node feature size: 6,805
+
+    Parameters
+    ----------
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of classes for each node.
+
+    Examples
+    --------
+    >>> data = CoauthorCSDataset()
+    >>> g = data[0]
+    >>> num_class = data.num_classes
+    >>> feat = g.ndata['feat']  # get node feature
+    >>> label = g.ndata['label']  # get node labels
+    """
+
+    def __init__(
+        self, raw_dir=None, force_reload=False, verbose=False, transform=None
+    ):
+        super(CoauthorCSDataset, self).__init__(
+            name="coauthor_cs",
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    @property
+    def num_classes(self):
+        """Number of classes.
+
+        Return
+        -------
+        int
+        """
+        return 15
+
+
+class CoauthorPhysicsDataset(GNNBenchmarkDataset):
+    r"""'Physics' part of the Coauthor dataset for node classification task.
+
+    Coauthor CS and Coauthor Physics are co-authorship graphs based on the Microsoft Academic Graph
+    from the KDD Cup 2016 challenge. Here, nodes are authors, that are connected by an edge if they
+    co-authored a paper; node features represent paper keywords for each author’s papers, and class
+    labels indicate most active fields of study for each author.
+
+    Reference: `<https://github.com/shchur/gnn-benchmark#datasets>`_
+
+    Statistics
+
+    - Nodes: 34,493
+    - Edges: 495,924 (note that the original dataset has 247,962 edges but DGL adds
+      the reverse edges and remove the duplicates, hence with a different number)
+    - Number of classes: 5
+    - Node feature size: 8,415
+
+    Parameters
+    ----------
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of classes for each node.
+
+    Examples
+    --------
+    >>> data = CoauthorPhysicsDataset()
+    >>> g = data[0]
+    >>> num_class = data.num_classes
+    >>> feat = g.ndata['feat']  # get node feature
+    >>> label = g.ndata['label']  # get node labels
+    """
+
+    def __init__(
+        self, raw_dir=None, force_reload=False, verbose=False, transform=None
+    ):
+        super(CoauthorPhysicsDataset, self).__init__(
+            name="coauthor_physics",
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    @property
+    def num_classes(self):
+        """Number of classes.
+
+        Return
+        -------
+        int
+        """
+        return 5
+
+
+class AmazonCoBuyComputerDataset(GNNBenchmarkDataset):
+    r"""'Computer' part of the AmazonCoBuy dataset for node classification task.
+
+    Amazon Computers and Amazon Photo are segments of the Amazon co-purchase graph [McAuley et al., 2015],
+    where nodes represent goods, edges indicate that two goods are frequently bought together, node
+    features are bag-of-words encoded product reviews, and class labels are given by the product category.
+
+    Reference: `<https://github.com/shchur/gnn-benchmark#datasets>`_
+
+    Statistics:
+
+    - Nodes: 13,752
+    - Edges: 491,722 (note that the original dataset has 245,778 edges but DGL adds
+      the reverse edges and remove the duplicates, hence with a different number)
+    - Number of classes: 10
+    - Node feature size: 767
+
+    Parameters
+    ----------
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of classes for each node.
+
+    Examples
+    --------
+    >>> data = AmazonCoBuyComputerDataset()
+    >>> g = data[0]
+    >>> num_class = data.num_classes
+    >>> feat = g.ndata['feat']  # get node feature
+    >>> label = g.ndata['label']  # get node labels
+    """
+
+    def __init__(
+        self, raw_dir=None, force_reload=False, verbose=False, transform=None
+    ):
+        super(AmazonCoBuyComputerDataset, self).__init__(
+            name="amazon_co_buy_computer",
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    @property
+    def num_classes(self):
+        """Number of classes.
+
+        Return
+        -------
+        int
+        """
+        return 10
+
+
+class AmazonCoBuyPhotoDataset(GNNBenchmarkDataset):
+    r"""AmazonCoBuy dataset for node classification task.
+
+    Amazon Computers and Amazon Photo are segments of the Amazon co-purchase graph [McAuley et al., 2015],
+    where nodes represent goods, edges indicate that two goods are frequently bought together, node
+    features are bag-of-words encoded product reviews, and class labels are given by the product category.
+
+    Reference: `<https://github.com/shchur/gnn-benchmark#datasets>`_
+
+    Statistics
+
+    - Nodes: 7,650
+    - Edges: 238,163 (note that the original dataset has 119,043 edges but DGL adds
+      the reverse edges and remove the duplicates, hence with a different number)
+    - Number of classes: 8
+    - Node feature size: 745
+
+    Parameters
+    ----------
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of classes for each node.
+
+    Examples
+    --------
+    >>> data = AmazonCoBuyPhotoDataset()
+    >>> g = data[0]
+    >>> num_class = data.num_classes
+    >>> feat = g.ndata['feat']  # get node feature
+    >>> label = g.ndata['label']  # get node labels
+    """
+
+    def __init__(
+        self, raw_dir=None, force_reload=False, verbose=False, transform=None
+    ):
+        super(AmazonCoBuyPhotoDataset, self).__init__(
+            name="amazon_co_buy_photo",
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    @property
+    def num_classes(self):
+        """Number of classes.
+
+        Return
+        -------
+        int
+        """
+        return 8
+
+
+class CoraFull(CoraFullDataset):
+    def __init__(self, **kwargs):
+        deprecate_class("CoraFull", "CoraFullDataset")
+        super(CoraFull, self).__init__(**kwargs)
+
+
+def AmazonCoBuy(name):
+    if name == "computers":
+        deprecate_class("AmazonCoBuy", "AmazonCoBuyComputerDataset")
+        return AmazonCoBuyComputerDataset()
+    elif name == "photo":
+        deprecate_class("AmazonCoBuy", "AmazonCoBuyPhotoDataset")
+        return AmazonCoBuyPhotoDataset()
+    else:
+        raise ValueError('Dataset name should be "computers" or "photo".')
+
+
+def Coauthor(name):
+    if name == "cs":
+        deprecate_class("Coauthor", "CoauthorCSDataset")
+        return CoauthorCSDataset()
+    elif name == "physics":
+        deprecate_class("Coauthor", "CoauthorPhysicsDataset")
+        return CoauthorPhysicsDataset()
+    else:
+        raise ValueError('Dataset name should be "cs" or "physics".')

python/dgl/data/graph_serialize.py

@@ -0,0 +1,272 @@
+"""For Graph Serialization"""
+from __future__ import absolute_import
+
+import os
+
+from .. import backend as F
+from .._ffi.function import _init_api
+from .._ffi.object import ObjectBase, register_object
+from ..base import dgl_warning, DGLError
+from ..heterograph import DGLGraph
+from .heterograph_serialize import save_heterographs
+
+_init_api("dgl.data.graph_serialize")
+
+__all__ = ["save_graphs", "load_graphs", "load_labels"]
+
+
+@register_object("graph_serialize.StorageMetaData")
+class StorageMetaData(ObjectBase):
+    """StorageMetaData Object
+    attributes available:
+      num_graph [int]: return numbers of graphs
+      nodes_num_list Value of NDArray: return number of nodes for each graph
+      edges_num_list Value of NDArray: return number of edges for each graph
+      labels [dict of backend tensors]: return dict of labels
+      graph_data [list of GraphData]: return list of GraphData Object
+    """
+
+
+def is_local_path(filepath):
+    return not (
+        filepath.startswith("hdfs://")
+        or filepath.startswith("viewfs://")
+        or filepath.startswith("s3://")
+    )
+
+
+def check_local_file_exists(filename):
+    if is_local_path(filename) and not os.path.exists(filename):
+        raise DGLError("File {} does not exist.".format(filename))
+
+
+@register_object("graph_serialize.GraphData")
+class GraphData(ObjectBase):
+    """GraphData Object"""
+
+    @staticmethod
+    def create(g):
+        """Create GraphData"""
+        # TODO(zihao): support serialize batched graph in the future.
+        assert (
+            g.batch_size == 1
+        ), "Batched DGLGraph is not supported for serialization"
+        ghandle = g._graph
+        if len(g.ndata) != 0:
+            node_tensors = dict()
+            for key, value in g.ndata.items():
+                node_tensors[key] = F.zerocopy_to_dgl_ndarray(value)
+        else:
+            node_tensors = None
+        if len(g.edata) != 0:
+            edge_tensors = dict()
+            for key, value in g.edata.items():
+                edge_tensors[key] = F.zerocopy_to_dgl_ndarray(value)
+        else:
+            edge_tensors = None
+        return _CAPI_MakeGraphData(ghandle, node_tensors, edge_tensors)
+
+    def get_graph(self):
+        """Get DGLGraph from GraphData"""
+        ghandle = _CAPI_GDataGraphHandle(self)
+        hgi = _CAPI_DGLAsHeteroGraph(ghandle)
+        g = DGLGraph(hgi, ["_U"], ["_E"])
+        node_tensors_items = _CAPI_GDataNodeTensors(self).items()
+        edge_tensors_items = _CAPI_GDataEdgeTensors(self).items()
+        for k, v in node_tensors_items:
+            g.ndata[k] = F.zerocopy_from_dgl_ndarray(v)
+        for k, v in edge_tensors_items:
+            g.edata[k] = F.zerocopy_from_dgl_ndarray(v)
+        return g
+
+
+def save_graphs(filename, g_list, labels=None, formats=None):
+    r"""Save graphs and optionally their labels to file.
+
+    Besides saving to local files, DGL supports writing the graphs directly
+    to S3 (by providing a ``"s3://..."`` path) or to HDFS (by providing
+    ``"hdfs://..."`` a path).
+
+    The function saves both the graph structure and node/edge features to file
+    in DGL's own binary format. For graph-level features, pass them via
+    the :attr:`labels` argument.
+
+    Parameters
+    ----------
+    filename : str
+        The file name to store the graphs and labels.
+    g_list: list
+        The graphs to be saved.
+    labels: dict[str, Tensor]
+        labels should be dict of tensors, with str as keys
+    formats: str or list[str]
+        Save graph in specified formats. It could be any combination of
+        ``coo``, ``csc`` and ``csr``. If not specified, save one format
+        only according to what format is available. If multiple formats
+        are available, selection priority from high to low is ``coo``,
+        ``csc``, ``csr``.
+
+    Examples
+    ----------
+    >>> import dgl
+    >>> import torch as th
+
+    Create :class:`DGLGraph` objects and initialize node
+    and edge features.
+
+    >>> g1 = dgl.graph(([0, 1, 2], [1, 2, 3]))
+    >>> g2 = dgl.graph(([0, 2], [2, 3]))
+    >>> g2.edata["e"] = th.ones(2, 4)
+
+    Save Graphs into file
+
+    >>> from dgl.data.utils import save_graphs
+    >>> graph_labels = {"glabel": th.tensor([0, 1])}
+    >>> save_graphs("./data.bin", [g1, g2], graph_labels)
+
+    See Also
+    --------
+    load_graphs
+    """
+    # if it is local file, do some sanity check
+    if is_local_path(filename):
+        if os.path.isdir(filename):
+            raise DGLError(
+                "Filename {} is an existing directory.".format(filename)
+            )
+        f_path = os.path.dirname(filename)
+        if f_path and not os.path.exists(f_path):
+            os.makedirs(f_path)
+    g_sample = g_list[0] if isinstance(g_list, list) else g_list
+    if type(g_sample) == DGLGraph:  # Doesn't support DGLGraph's derived class
+        save_heterographs(filename, g_list, labels, formats)
+    else:
+        raise DGLError(
+            "Invalid argument g_list. Must be a DGLGraph or a list of DGLGraphs."
+        )
+
+
+def load_graphs(filename, idx_list=None):
+    """Load graphs and optionally their labels from file saved by :func:`save_graphs`.
+
+    Besides loading from local files, DGL supports loading the graphs directly
+    from S3 (by providing a ``"s3://..."`` path) or from HDFS (by providing
+    ``"hdfs://..."`` a path).
+
+    Parameters
+    ----------
+    filename: str
+        The file name to load graphs from.
+    idx_list: list[int], optional
+        The indices of the graphs to be loaded if the file contains multiple graphs.
+        Default is loading all the graphs stored in the file.
+
+    Returns
+    --------
+    graph_list: list[DGLGraph]
+        The loaded graphs.
+    labels: dict[str, Tensor]
+        The graph labels stored in file. If no label is stored, the dictionary is empty.
+        Regardless of whether the ``idx_list`` argument is given or not,
+        the returned dictionary always contains the labels of all the graphs.
+
+    Examples
+    ----------
+    Following the example in :func:`save_graphs`.
+
+    >>> from dgl.data.utils import load_graphs
+    >>> glist, label_dict = load_graphs("./data.bin") # glist will be [g1, g2]
+    >>> glist, label_dict = load_graphs("./data.bin", [0]) # glist will be [g1]
+
+    See Also
+    --------
+    save_graphs
+    """
+    # if it is local file, do some sanity check
+    check_local_file_exists(filename)
+    version = _CAPI_GetFileVersion(filename)
+    if version == 1:
+        dgl_warning(
+            "You are loading a graph file saved by old version of dgl.  \
+            Please consider saving it again with the current format."
+        )
+        return load_graph_v1(filename, idx_list)
+    elif version == 2:
+        return load_graph_v2(filename, idx_list)
+    else:
+        raise DGLError("Invalid DGL Version Number.")
+
+
+def load_graph_v2(filename, idx_list=None):
+    """Internal functions for loading DGLGraphs."""
+    if idx_list is None:
+        idx_list = []
+    assert isinstance(idx_list, list)
+    heterograph_list = _CAPI_LoadGraphFiles_V2(filename, idx_list)
+    label_dict = load_labels_v2(filename)
+    return [gdata.get_graph() for gdata in heterograph_list], label_dict
+
+
+def load_graph_v1(filename, idx_list=None):
+    """ "Internal functions for loading DGLGraphs (V0)."""
+    if idx_list is None:
+        idx_list = []
+    assert isinstance(idx_list, list)
+    metadata = _CAPI_LoadGraphFiles_V1(filename, idx_list, False)
+    label_dict = {}
+    for k, v in metadata.labels.items():
+        label_dict[k] = F.zerocopy_from_dgl_ndarray(v)
+    return [gdata.get_graph() for gdata in metadata.graph_data], label_dict
+
+
+def load_labels(filename):
+    """
+    Load label dict from file
+
+    Parameters
+    ----------
+    filename: str
+        filename to load DGLGraphs
+
+    Returns
+    ----------
+    labels: dict
+        dict of labels stored in file (empty dict returned if no
+        label stored)
+
+    Examples
+    ----------
+    Following the example in save_graphs.
+
+    >>> from dgl.data.utils import load_labels
+    >>> label_dict = load_graphs("./data.bin")
+
+    """
+    # if it is local file, do some sanity check
+    check_local_file_exists(filename)
+
+    version = _CAPI_GetFileVersion(filename)
+    if version == 1:
+        return load_labels_v1(filename)
+    elif version == 2:
+        return load_labels_v2(filename)
+    else:
+        raise Exception("Invalid DGL Version Number")
+
+
+def load_labels_v2(filename):
+    """Internal functions for loading labels from V2 format"""
+    label_dict = {}
+    nd_dict = _CAPI_LoadLabels_V2(filename)
+    for k, v in nd_dict.items():
+        label_dict[k] = F.zerocopy_from_dgl_ndarray(v)
+    return label_dict
+
+
+def load_labels_v1(filename):
+    """Internal functions for loading labels from V1 format"""
+    metadata = _CAPI_LoadGraphFiles_V1(filename, [], True)
+    label_dict = {}
+    for k, v in metadata.labels.items():
+        label_dict[k] = F.zerocopy_from_dgl_ndarray(v)
+    return label_dict

python/dgl/data/heterograph_serialize.py

@@ -0,0 +1,79 @@
+"""For HeteroGraph Serialization"""
+from __future__ import absolute_import
+
+from .. import backend as F
+from .._ffi.function import _init_api
+from .._ffi.object import ObjectBase, register_object
+from ..container import convert_to_strmap
+from ..frame import Frame
+from ..heterograph import DGLGraph
+
+_init_api("dgl.data.heterograph_serialize")
+
+
+def tensor_dict_to_ndarray_dict(tensor_dict):
+    """Convert dict[str, tensor] to StrMap[NDArray]"""
+    ndarray_dict = {}
+    for key, value in tensor_dict.items():
+        ndarray_dict[key] = F.zerocopy_to_dgl_ndarray(value)
+    return convert_to_strmap(ndarray_dict)
+
+
+def save_heterographs(filename, g_list, labels, formats):
+    """Save heterographs into file"""
+    if labels is None:
+        labels = {}
+    if isinstance(g_list, DGLGraph):
+        g_list = [g_list]
+    assert all(
+        [type(g) == DGLGraph for g in g_list]
+    ), "Invalid DGLGraph in g_list argument"
+    gdata_list = [HeteroGraphData.create(g) for g in g_list]
+    if formats is None:
+        formats = []
+    elif isinstance(formats, str):
+        formats = [formats]
+    _CAPI_SaveHeteroGraphData(
+        filename, gdata_list, tensor_dict_to_ndarray_dict(labels), formats
+    )
+
+
+@register_object("heterograph_serialize.HeteroGraphData")
+class HeteroGraphData(ObjectBase):
+    """Object to hold the data to be stored for DGLGraph"""
+
+    @staticmethod
+    def create(g):
+        edata_list = []
+        ndata_list = []
+        for etype in g.canonical_etypes:
+            edata_list.append(tensor_dict_to_ndarray_dict(g.edges[etype].data))
+        for ntype in g.ntypes:
+            ndata_list.append(tensor_dict_to_ndarray_dict(g.nodes[ntype].data))
+        return _CAPI_MakeHeteroGraphData(
+            g._graph, ndata_list, edata_list, g.ntypes, g.etypes
+        )
+
+    def get_graph(self):
+        ntensor_list = list(_CAPI_GetNDataFromHeteroGraphData(self))
+        etensor_list = list(_CAPI_GetEDataFromHeteroGraphData(self))
+        ntype_names = list(_CAPI_GetNtypesFromHeteroGraphData(self))
+        etype_names = list(_CAPI_GetEtypesFromHeteroGraphData(self))
+        gidx = _CAPI_GetGindexFromHeteroGraphData(self)
+        nframes = []
+        eframes = []
+        for ntid, ntensor in enumerate(ntensor_list):
+            ndict = {
+                ntensor[i]: F.zerocopy_from_dgl_ndarray(ntensor[i + 1])
+                for i in range(0, len(ntensor), 2)
+            }
+            nframes.append(Frame(ndict, num_rows=gidx.num_nodes(ntid)))
+
+        for etid, etensor in enumerate(etensor_list):
+            edict = {
+                etensor[i]: F.zerocopy_from_dgl_ndarray(etensor[i + 1])
+                for i in range(0, len(etensor), 2)
+            }
+            eframes.append(Frame(edict, num_rows=gidx.num_edges(etid)))
+
+        return DGLGraph(gidx, ntype_names, etype_names, nframes, eframes)

python/dgl/data/heterophilous_graphs.py

@@ -0,0 +1,456 @@
+"""
+Datasets introduced in the 'A Critical Look at the Evaluation of GNNs under Heterophily: Are We
+Really Making Progress? <https://arxiv.org/abs/2302.11640>'__ paper.
+"""
+import os
+
+import numpy as np
+
+from ..convert import graph
+from ..transforms.functional import to_bidirected
+from .dgl_dataset import DGLBuiltinDataset
+from .utils import download
+
+
+class HeterophilousGraphDataset(DGLBuiltinDataset):
+    r"""Datasets introduced in the 'A Critical Look at the Evaluation of GNNs under Heterophily:
+    Are We Really Making Progress? <https://arxiv.org/abs/2302.11640>'__ paper.
+
+    Parameters
+    ----------
+    name : str
+        Name of the dataset. One of 'roman-empire', 'amazon-ratings', 'minesweeper', 'tolokers',
+        'questions'.
+    raw_dir : str
+        Raw file directory to store the processed data.
+    force_reload : bool
+        Whether to re-download the data source.
+    verbose : bool
+        Whether to print progress information.
+    transform : callable
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+    """
+
+    def __init__(
+        self,
+        name,
+        raw_dir=None,
+        force_reload=False,
+        verbose=True,
+        transform=None,
+    ):
+        name = name.lower().replace("-", "_")
+        url = f"https://github.com/yandex-research/heterophilous-graphs/raw/main/data/{name}.npz"
+        super(HeterophilousGraphDataset, self).__init__(
+            name=name,
+            url=url,
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def download(self):
+        download(
+            url=self.url, path=os.path.join(self.raw_path, f"{self.name}.npz")
+        )
+
+    def process(self):
+        """Load and process the data."""
+        try:
+            import torch
+        except ImportError:
+            raise ModuleNotFoundError(
+                "This dataset requires PyTorch to be the backend."
+            )
+
+        data = np.load(os.path.join(self.raw_path, f"{self.name}.npz"))
+        src = torch.from_numpy(data["edges"][:, 0])
+        dst = torch.from_numpy(data["edges"][:, 1])
+        features = torch.from_numpy(data["node_features"])
+        labels = torch.from_numpy(data["node_labels"])
+        train_masks = torch.from_numpy(data["train_masks"].T)
+        val_masks = torch.from_numpy(data["val_masks"].T)
+        test_masks = torch.from_numpy(data["test_masks"].T)
+        num_nodes = len(labels)
+        num_classes = len(labels.unique())
+
+        self._num_classes = num_classes
+
+        self._g = to_bidirected(graph((src, dst), num_nodes=num_nodes))
+        self._g.ndata["feat"] = features
+        self._g.ndata["label"] = labels
+        self._g.ndata["train_mask"] = train_masks
+        self._g.ndata["val_mask"] = val_masks
+        self._g.ndata["test_mask"] = test_masks
+
+    def has_cache(self):
+        return os.path.exists(self.raw_path)
+
+    def load(self):
+        self.process()
+
+    def __getitem__(self, idx):
+        assert idx == 0, "This dataset has only one graph."
+        if self._transform is None:
+            return self._g
+        else:
+            return self._transform(self._g)
+
+    def __len__(self):
+        return 1
+
+    @property
+    def num_classes(self):
+        return self._num_classes
+
+
+class RomanEmpireDataset(HeterophilousGraphDataset):
+    r"""Roman-empire dataset from the 'A Critical Look at the Evaluation of GNNs under Heterophily:
+    Are We Really Making Progress? <https://arxiv.org/abs/2302.11640>'__ paper.
+
+    This dataset is based on the Roman Empire article from English Wikipedia, which was selected
+    since it is one of the longest articles on Wikipedia. Each node in the graph corresponds to one
+    (non-unique) word in the text. Thus, the number of nodes in the graph is equal to the article’s
+    length. Two words are connected with an edge if at least one of the following two conditions
+    holds: either these words follow each other in the text, or these words are connected in the
+    dependency tree of the sentence (one word depends on the other). Thus, the graph is a chain
+    graph with additional shortcut edges corresponding to syntactic dependencies between words. The
+    class of a node is its syntactic role (17 most frequent roles were selected as unique classes
+    and all the other roles were grouped into the 18th class). Node features are word embeddings.
+
+    Statistics:
+
+    - Nodes: 22662
+    - Edges: 65854
+    - Classes: 18
+    - Node features: 300
+    - 10 train/val/test splits
+
+    Parameters
+    ----------
+    raw_dir : str, optional
+        Raw file directory to store the processed data. Default: ~/.dgl/
+    force_reload : bool, optional
+        Whether to re-download the data source. Default: False
+    verbose : bool, optional
+        Whether to print progress information. Default: True
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access. Default: None
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of node classes
+
+
+    Examples
+    --------
+
+    >>> from dgl.data import RomanEmpireDataset
+    >>> dataset = RomanEmpireDataset()
+    >>> g = dataset[0]
+    >>> num_classes = dataset.num_classes
+
+    >>> # get node features
+    >>> feat = g.ndata["feat"]
+
+    >>> # get the first data split
+    >>> train_mask = g.ndata["train_mask"][:, 0]
+    >>> val_mask = g.ndata["val_mask"][:, 0]
+    >>> test_mask = g.ndata["test_mask"][:, 0]
+
+    >>> # get labels
+    >>> label = g.ndata['label']
+    """
+
+    def __init__(
+        self, raw_dir=None, force_reload=False, verbose=True, transform=None
+    ):
+        super(RomanEmpireDataset, self).__init__(
+            name="roman-empire",
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+
+class AmazonRatingsDataset(HeterophilousGraphDataset):
+    r"""Amazon-ratings dataset from the 'A Critical Look at the Evaluation of GNNs under
+    Heterophily: Are We Really Making Progress? <https://arxiv.org/abs/2302.11640>'__ paper.
+
+    This dataset is based on the Amazon product co-purchasing data. Nodes are products (books, music
+    CDs, DVDs, VHS video tapes), and edges connect products that are frequently bought together. The
+    task is to predict the average rating given to a product by reviewers. All possible rating
+    values were grouped into five classes. Node features are the mean of word embeddings for words
+    in the product description.
+
+    Statistics:
+
+    - Nodes: 24492
+    - Edges: 186100
+    - Classes: 5
+    - Node features: 300
+    - 10 train/val/test splits
+
+    Parameters
+    ----------
+    raw_dir : str, optional
+        Raw file directory to store the processed data. Default: ~/.dgl/
+    force_reload : bool, optional
+        Whether to re-download the data source. Default: False
+    verbose : bool, optional
+        Whether to print progress information. Default: True
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access. Default: None
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of node classes
+
+
+    Examples
+    --------
+
+    >>> from dgl.data import AmazonRatingsDataset
+    >>> dataset = AmazonRatingsDataset()
+    >>> g = dataset[0]
+    >>> num_classes = dataset.num_classes
+
+    >>> # get node features
+    >>> feat = g.ndata["feat"]
+
+    >>> # get the first data split
+    >>> train_mask = g.ndata["train_mask"][:, 0]
+    >>> val_mask = g.ndata["val_mask"][:, 0]
+    >>> test_mask = g.ndata["test_mask"][:, 0]
+
+    >>> # get labels
+    >>> label = g.ndata['label']
+    """
+
+    def __init__(
+        self, raw_dir=None, force_reload=False, verbose=True, transform=None
+    ):
+        super(AmazonRatingsDataset, self).__init__(
+            name="amazon-ratings",
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+
+class MinesweeperDataset(HeterophilousGraphDataset):
+    r"""Minesweeper dataset from the 'A Critical Look at the Evaluation of GNNs under Heterophily:
+    Are We Really Making Progress? <https://arxiv.org/abs/2302.11640>'__ paper.
+
+    This dataset is inspired by the Minesweeper game. The graph is a regular 100x100 grid where each
+    node (cell) is connected to eight neighboring nodes (with the exception of nodes at the edge of
+    the grid, which have fewer neighbors). 20% of the nodes are randomly selected as mines. The task
+    is to predict which nodes are mines. The node features are one-hot-encoded numbers of
+    neighboring mines. However, for randomly selected 50% of the nodes, the features are unknown,
+    which is indicated by a separate binary feature.
+
+    Statistics:
+
+    - Nodes: 10000
+    - Edges: 78804
+    - Classes: 2
+    - Node features: 7
+    - 10 train/val/test splits
+
+    Parameters
+    ----------
+    raw_dir : str, optional
+        Raw file directory to store the processed data. Default: ~/.dgl/
+    force_reload : bool, optional
+        Whether to re-download the data source. Default: False
+    verbose : bool, optional
+        Whether to print progress information. Default: True
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access. Default: None
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of node classes
+
+
+    Examples
+    --------
+
+    >>> from dgl.data import MinesweeperDataset
+    >>> dataset = MinesweeperDataset()
+    >>> g = dataset[0]
+    >>> num_classes = dataset.num_classes
+
+    >>> # get node features
+    >>> feat = g.ndata["feat"]
+
+    >>> # get the first data split
+    >>> train_mask = g.ndata["train_mask"][:, 0]
+    >>> val_mask = g.ndata["val_mask"][:, 0]
+    >>> test_mask = g.ndata["test_mask"][:, 0]
+
+    >>> # get labels
+    >>> label = g.ndata['label']
+    """
+
+    def __init__(
+        self, raw_dir=None, force_reload=False, verbose=True, transform=None
+    ):
+        super(MinesweeperDataset, self).__init__(
+            name="minesweeper",
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+
+class TolokersDataset(HeterophilousGraphDataset):
+    r"""Tolokers dataset from the 'A Critical Look at the Evaluation of GNNs under Heterophily:
+    Are We Really Making Progress? <https://arxiv.org/abs/2302.11640>'__ paper.
+
+    This dataset is based on data from the Toloka crowdsourcing platform. The nodes represent
+    tolokers (workers). An edge connects two tolokers if they have worked on the same task. The goal
+    is to predict which tolokers have been banned in one of the projects. Node features are based on
+    the worker’s profile information and task performance statistics.
+
+    Statistics:
+
+    - Nodes: 11758
+    - Edges: 1038000
+    - Classes: 2
+    - Node features: 10
+    - 10 train/val/test splits
+
+    Parameters
+    ----------
+    raw_dir : str, optional
+        Raw file directory to store the processed data. Default: ~/.dgl/
+    force_reload : bool, optional
+        Whether to re-download the data source. Default: False
+    verbose : bool, optional
+        Whether to print progress information. Default: True
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access. Default: None
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of node classes
+
+
+    Examples
+    --------
+
+    >>> from dgl.data import TolokersDataset
+    >>> dataset = TolokersDataset()
+    >>> g = dataset[0]
+    >>> num_classes = dataset.num_classes
+
+    >>> # get node features
+    >>> feat = g.ndata["feat"]
+
+    >>> # get the first data split
+    >>> train_mask = g.ndata["train_mask"][:, 0]
+    >>> val_mask = g.ndata["val_mask"][:, 0]
+    >>> test_mask = g.ndata["test_mask"][:, 0]
+
+    >>> # get labels
+    >>> label = g.ndata['label']
+    """
+
+    def __init__(
+        self, raw_dir=None, force_reload=False, verbose=True, transform=None
+    ):
+        super(TolokersDataset, self).__init__(
+            name="tolokers",
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+
+class QuestionsDataset(HeterophilousGraphDataset):
+    r"""Questions dataset from the 'A Critical Look at the Evaluation of GNNs under Heterophily:
+    Are We Really Making Progress? <https://arxiv.org/abs/2302.11640>'__ paper.
+
+    This dataset is based on data from the question-answering website Yandex Q. Nodes are users, and
+    an edge connects two nodes if one user answered the other user’s question. The task is to
+    predict which users remained active on the website (were not deleted or blocked). Node features
+    are the mean of word embeddings for words in the user description. Users that do not have
+    description are indicated by a separate binary feature.
+
+    Statistics:
+
+    - Nodes: 48921
+    - Edges: 307080
+    - Classes: 2
+    - Node features: 301
+    - 10 train/val/test splits
+
+    Parameters
+    ----------
+    raw_dir : str, optional
+        Raw file directory to store the processed data. Default: ~/.dgl/
+    force_reload : bool, optional
+        Whether to re-download the data source. Default: False
+    verbose : bool, optional
+        Whether to print progress information. Default: True
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access. Default: None
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of node classes
+
+
+    Examples
+    --------
+
+    >>> from dgl.data import QuestionsDataset
+    >>> dataset = QuestionsDataset()
+    >>> g = dataset[0]
+    >>> num_classes = dataset.num_classes
+
+    >>> # get node features
+    >>> feat = g.ndata["feat"]
+
+    >>> # get the first data split
+    >>> train_mask = g.ndata["train_mask"][:, 0]
+    >>> val_mask = g.ndata["val_mask"][:, 0]
+    >>> test_mask = g.ndata["test_mask"][:, 0]
+
+    >>> # get labels
+    >>> label = g.ndata['label']
+    """
+
+    def __init__(
+        self, raw_dir=None, force_reload=False, verbose=True, transform=None
+    ):
+        super(QuestionsDataset, self).__init__(
+            name="questions",
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )

python/dgl/data/icews18.py

@@ -0,0 +1,172 @@
+"""ICEWS18 dataset for temporal graph"""
+import os
+
+import numpy as np
+
+from .. import backend as F
+from ..convert import graph as dgl_graph
+from .dgl_dataset import DGLBuiltinDataset
+from .utils import _get_dgl_url, load_graphs, loadtxt, save_graphs
+
+
+class ICEWS18Dataset(DGLBuiltinDataset):
+    r"""ICEWS18 dataset for temporal graph
+
+    Integrated Crisis Early Warning System (ICEWS18)
+
+    Event data consists of coded interactions between socio-political
+    actors (i.e., cooperative or hostile actions between individuals,
+    groups, sectors and nation states). This Dataset consists of events
+    from 1/1/2018 to 10/31/2018 (24 hours time granularity).
+
+    Reference:
+
+        - `Recurrent Event Network for Reasoning over Temporal Knowledge Graphs <https://arxiv.org/abs/1904.05530>`_
+        - `ICEWS Coded Event Data <https://dataverse.harvard.edu/dataset.xhtml?persistentId=doi:10.7910/DVN/28075>`_
+
+    Statistics：
+
+    - Train examples: 240
+    - Valid examples: 30
+    - Test examples: 34
+    - Nodes per graph: 23033
+
+    Parameters
+    ----------
+    mode: str
+        Load train/valid/test data. Has to be one of ['train', 'valid', 'test']
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    -------
+    is_temporal : bool
+        Is the dataset contains temporal graphs
+
+    Examples
+    --------
+    >>> # get train, valid, test set
+    >>> train_data = ICEWS18Dataset()
+    >>> valid_data = ICEWS18Dataset(mode='valid')
+    >>> test_data = ICEWS18Dataset(mode='test')
+    >>>
+    >>> train_size = len(train_data)
+    >>> for g in train_data:
+    ....    e_feat = g.edata['rel_type']
+    ....    # your code here
+    ....
+    >>>
+    """
+
+    def __init__(
+        self,
+        mode="train",
+        raw_dir=None,
+        force_reload=False,
+        verbose=False,
+        transform=None,
+    ):
+        mode = mode.lower()
+        assert mode in ["train", "valid", "test"], "Mode not valid"
+        self.mode = mode
+        _url = _get_dgl_url("dataset/icews18.zip")
+        super(ICEWS18Dataset, self).__init__(
+            name="ICEWS18",
+            url=_url,
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        data = loadtxt(
+            os.path.join(self.save_path, "{}.txt".format(self.mode)),
+            delimiter="\t",
+        ).astype(np.int64)
+        num_nodes = 23033
+        # The source code is not released, but the paper indicates there're
+        # totally 137 samples. The cutoff below has exactly 137 samples.
+        time_index = np.floor(data[:, 3] / 24).astype(np.int64)
+        start_time = time_index[time_index != -1].min()
+        end_time = time_index.max()
+        self._graphs = []
+        for i in range(start_time, end_time + 1):
+            row_mask = time_index <= i
+            edges = data[row_mask][:, [0, 2]]
+            rate = data[row_mask][:, 1]
+            g = dgl_graph((edges[:, 0], edges[:, 1]))
+            g.edata["rel_type"] = F.tensor(
+                rate.reshape(-1, 1), dtype=F.data_type_dict["int64"]
+            )
+            self._graphs.append(g)
+
+    def has_cache(self):
+        graph_path = os.path.join(
+            self.save_path, "{}_dgl_graph.bin".format(self.mode)
+        )
+        return os.path.exists(graph_path)
+
+    def save(self):
+        graph_path = os.path.join(
+            self.save_path, "{}_dgl_graph.bin".format(self.mode)
+        )
+        save_graphs(graph_path, self._graphs)
+
+    def load(self):
+        graph_path = os.path.join(
+            self.save_path, "{}_dgl_graph.bin".format(self.mode)
+        )
+        self._graphs = load_graphs(graph_path)[0]
+
+    def __getitem__(self, idx):
+        r"""Get graph by index
+
+        Parameters
+        ----------
+        idx : int
+            Item index
+
+        Returns
+        -------
+        :class:`dgl.DGLGraph`
+
+            The graph contains:
+
+            - ``edata['rel_type']``: edge type
+        """
+        if self._transform is None:
+            return self._graphs[idx]
+        else:
+            return self._transform(self._graphs[idx])
+
+    def __len__(self):
+        r"""Number of graphs in the dataset.
+
+        Return
+        -------
+        int
+        """
+        return len(self._graphs)
+
+    @property
+    def is_temporal(self):
+        r"""Is the dataset contains temporal graphs
+
+        Returns
+        -------
+        bool
+        """
+        return True
+
+
+ICEWS18 = ICEWS18Dataset

python/dgl/data/karate.py

@@ -0,0 +1,98 @@
+"""KarateClub Dataset
+"""
+import networkx as nx
+import numpy as np
+
+from .. import backend as F
+from ..convert import from_networkx
+from .dgl_dataset import DGLDataset
+from .utils import deprecate_property
+
+__all__ = ["KarateClubDataset", "KarateClub"]
+
+
+class KarateClubDataset(DGLDataset):
+    r"""Karate Club dataset for Node Classification
+
+    Zachary's karate club is a social network of a university
+    karate club, described in the paper "An Information Flow
+    Model for Conflict and Fission in Small Groups" by Wayne W. Zachary.
+    The network became a popular example of community structure in
+    networks after its use by Michelle Girvan and Mark Newman in 2002.
+    Official website: `<http://konect.cc/networks/ucidata-zachary/>`_
+
+    Karate Club dataset statistics:
+
+    - Nodes: 34
+    - Edges: 156
+    - Number of Classes: 2
+
+    Parameters
+    ----------
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of node classes
+
+    Examples
+    --------
+    >>> dataset = KarateClubDataset()
+    >>> num_classes = dataset.num_classes
+    >>> g = dataset[0]
+    >>> labels = g.ndata['label']
+    """
+
+    def __init__(self, transform=None):
+        super(KarateClubDataset, self).__init__(
+            name="karate_club", transform=transform
+        )
+
+    def process(self):
+        kc_graph = nx.karate_club_graph()
+        label = np.asarray(
+            [kc_graph.nodes[i]["club"] != "Mr. Hi" for i in kc_graph.nodes]
+        ).astype(np.int64)
+        label = F.tensor(label)
+        g = from_networkx(kc_graph)
+        g.ndata["label"] = label
+        self._graph = g
+        self._data = [g]
+
+    @property
+    def num_classes(self):
+        """Number of classes."""
+        return 2
+
+    def __getitem__(self, idx):
+        r"""Get graph object
+
+        Parameters
+        ----------
+        idx : int
+            Item index, KarateClubDataset has only one graph object
+
+        Returns
+        -------
+        :class:`dgl.DGLGraph`
+
+            graph structure and labels.
+
+            - ``ndata['label']``: ground truth labels
+        """
+        assert idx == 0, "This dataset has only one graph"
+        if self._transform is None:
+            return self._graph
+        else:
+            return self._transform(self._graph)
+
+    def __len__(self):
+        r"""The number of graphs in the dataset."""
+        return 1
+
+
+KarateClub = KarateClubDataset

python/dgl/data/knowledge_graph.py

@@ -0,0 +1,779 @@
+from __future__ import absolute_import
+
+import os, sys
+import pickle as pkl
+
+import networkx as nx
+
+import numpy as np
+import scipy.sparse as sp
+
+from .. import backend as F
+from ..convert import graph as dgl_graph
+from ..utils import retry_method_with_fix
+
+from .dgl_dataset import DGLBuiltinDataset
+from .utils import (
+    _get_dgl_url,
+    deprecate_function,
+    deprecate_property,
+    download,
+    extract_archive,
+    generate_mask_tensor,
+    get_download_dir,
+    load_graphs,
+    load_info,
+    makedirs,
+    save_graphs,
+    save_info,
+)
+
+
+class KnowledgeGraphDataset(DGLBuiltinDataset):
+    """KnowledgeGraph link prediction dataset
+
+    The dataset contains a graph depicting the connectivity of a knowledge
+    base. Currently, the knowledge bases from the
+    `RGCN paper <https://arxiv.org/pdf/1703.06103.pdf>`_ supported are
+    FB15k-237, FB15k, wn18
+
+    Parameters
+    -----------
+    name : str
+        Name can be 'FB15k-237', 'FB15k' or 'wn18'.
+    reverse : bool
+        Whether add reverse edges. Default: True.
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+    """
+
+    def __init__(
+        self,
+        name,
+        reverse=True,
+        raw_dir=None,
+        force_reload=False,
+        verbose=True,
+        transform=None,
+    ):
+        self._name = name
+        self.reverse = reverse
+        url = _get_dgl_url("dataset/") + "{}.tgz".format(name)
+        super(KnowledgeGraphDataset, self).__init__(
+            name,
+            url=url,
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def download(self):
+        r"""Automatically download data and extract it."""
+        tgz_path = os.path.join(self.raw_dir, self.name + ".tgz")
+        download(self.url, path=tgz_path)
+        extract_archive(tgz_path, self.raw_path)
+
+    def process(self):
+        """
+        The original knowledge base is stored in triplets.
+        This function will parse these triplets and build the DGLGraph.
+        """
+        root_path = self.raw_path
+        entity_path = os.path.join(root_path, "entities.dict")
+        relation_path = os.path.join(root_path, "relations.dict")
+        train_path = os.path.join(root_path, "train.txt")
+        valid_path = os.path.join(root_path, "valid.txt")
+        test_path = os.path.join(root_path, "test.txt")
+        entity_dict = _read_dictionary(entity_path)
+        relation_dict = _read_dictionary(relation_path)
+        train = np.asarray(
+            _read_triplets_as_list(train_path, entity_dict, relation_dict)
+        )
+        valid = np.asarray(
+            _read_triplets_as_list(valid_path, entity_dict, relation_dict)
+        )
+        test = np.asarray(
+            _read_triplets_as_list(test_path, entity_dict, relation_dict)
+        )
+        num_nodes = len(entity_dict)
+        num_rels = len(relation_dict)
+        if self.verbose:
+            print("# entities: {}".format(num_nodes))
+            print("# relations: {}".format(num_rels))
+            print("# training edges: {}".format(train.shape[0]))
+            print("# validation edges: {}".format(valid.shape[0]))
+            print("# testing edges: {}".format(test.shape[0]))
+
+        # for compatability
+        self._train = train
+        self._valid = valid
+        self._test = test
+
+        self._num_nodes = num_nodes
+        self._num_rels = num_rels
+        # build graph
+        g, data = build_knowledge_graph(
+            num_nodes, num_rels, train, valid, test, reverse=self.reverse
+        )
+        (
+            etype,
+            ntype,
+            train_edge_mask,
+            valid_edge_mask,
+            test_edge_mask,
+            train_mask,
+            val_mask,
+            test_mask,
+        ) = data
+        g.edata["train_edge_mask"] = train_edge_mask
+        g.edata["valid_edge_mask"] = valid_edge_mask
+        g.edata["test_edge_mask"] = test_edge_mask
+        g.edata["train_mask"] = train_mask
+        g.edata["val_mask"] = val_mask
+        g.edata["test_mask"] = test_mask
+        g.edata["etype"] = etype
+        g.ndata["ntype"] = ntype
+        self._g = g
+
+    @property
+    def graph_path(self):
+        return os.path.join(self.save_path, self.save_name + ".bin")
+
+    @property
+    def info_path(self):
+        return os.path.join(self.save_path, self.save_name + ".pkl")
+
+    def has_cache(self):
+        if os.path.exists(self.graph_path) and os.path.exists(self.info_path):
+            return True
+
+        return False
+
+    def __getitem__(self, idx):
+        assert idx == 0, "This dataset has only one graph"
+        if self._transform is None:
+            return self._g
+        else:
+            return self._transform(self._g)
+
+    def __len__(self):
+        return 1
+
+    def save(self):
+        """save the graph list and the labels"""
+        save_graphs(str(self.graph_path), self._g)
+        save_info(
+            str(self.info_path),
+            {"num_nodes": self.num_nodes, "num_rels": self.num_rels},
+        )
+
+    def load(self):
+        graphs, _ = load_graphs(str(self.graph_path))
+
+        info = load_info(str(self.info_path))
+        self._num_nodes = info["num_nodes"]
+        self._num_rels = info["num_rels"]
+        self._g = graphs[0]
+        train_mask = self._g.edata["train_edge_mask"].numpy()
+        val_mask = self._g.edata["valid_edge_mask"].numpy()
+        test_mask = self._g.edata["test_edge_mask"].numpy()
+
+        # convert mask tensor into bool tensor if possible
+        self._g.edata["train_edge_mask"] = generate_mask_tensor(
+            self._g.edata["train_edge_mask"].numpy()
+        )
+        self._g.edata["valid_edge_mask"] = generate_mask_tensor(
+            self._g.edata["valid_edge_mask"].numpy()
+        )
+        self._g.edata["test_edge_mask"] = generate_mask_tensor(
+            self._g.edata["test_edge_mask"].numpy()
+        )
+        self._g.edata["train_mask"] = generate_mask_tensor(
+            self._g.edata["train_mask"].numpy()
+        )
+        self._g.edata["val_mask"] = generate_mask_tensor(
+            self._g.edata["val_mask"].numpy()
+        )
+        self._g.edata["test_mask"] = generate_mask_tensor(
+            self._g.edata["test_mask"].numpy()
+        )
+
+        # for compatability (with 0.4.x) generate train_idx, valid_idx and test_idx
+        etype = self._g.edata["etype"].numpy()
+        self._etype = etype
+        u, v = self._g.all_edges(form="uv")
+        u = u.numpy()
+        v = v.numpy()
+        train_idx = np.nonzero(train_mask == 1)
+        self._train = np.column_stack(
+            (u[train_idx], etype[train_idx], v[train_idx])
+        )
+        valid_idx = np.nonzero(val_mask == 1)
+        self._valid = np.column_stack(
+            (u[valid_idx], etype[valid_idx], v[valid_idx])
+        )
+        test_idx = np.nonzero(test_mask == 1)
+        self._test = np.column_stack(
+            (u[test_idx], etype[test_idx], v[test_idx])
+        )
+
+        if self.verbose:
+            print("# entities: {}".format(self.num_nodes))
+            print("# relations: {}".format(self.num_rels))
+            print("# training edges: {}".format(self._train.shape[0]))
+            print("# validation edges: {}".format(self._valid.shape[0]))
+            print("# testing edges: {}".format(self._test.shape[0]))
+
+    @property
+    def num_nodes(self):
+        return self._num_nodes
+
+    @property
+    def num_rels(self):
+        return self._num_rels
+
+    @property
+    def save_name(self):
+        return self.name + "_dgl_graph"
+
+
+def _read_dictionary(filename):
+    d = {}
+    with open(filename, "r+") as f:
+        for line in f:
+            line = line.strip().split("\t")
+            d[line[1]] = int(line[0])
+    return d
+
+
+def _read_triplets(filename):
+    with open(filename, "r+") as f:
+        for line in f:
+            processed_line = line.strip().split("\t")
+            yield processed_line
+
+
+def _read_triplets_as_list(filename, entity_dict, relation_dict):
+    l = []
+    for triplet in _read_triplets(filename):
+        s = entity_dict[triplet[0]]
+        r = relation_dict[triplet[1]]
+        o = entity_dict[triplet[2]]
+        l.append([s, r, o])
+    return l
+
+
+def build_knowledge_graph(
+    num_nodes, num_rels, train, valid, test, reverse=True
+):
+    """Create a DGL Homogeneous graph with heterograph info stored as node or edge features."""
+    src = []
+    rel = []
+    dst = []
+    raw_subg = {}
+    raw_subg_eset = {}
+    raw_subg_etype = {}
+    raw_reverse_sugb = {}
+    raw_reverse_subg_eset = {}
+    raw_reverse_subg_etype = {}
+
+    # here there is noly one node type
+    s_type = "node"
+    d_type = "node"
+
+    def add_edge(s, r, d, reverse, edge_set):
+        r_type = str(r)
+        e_type = (s_type, r_type, d_type)
+        if raw_subg.get(e_type, None) is None:
+            raw_subg[e_type] = ([], [])
+            raw_subg_eset[e_type] = []
+            raw_subg_etype[e_type] = []
+        raw_subg[e_type][0].append(s)
+        raw_subg[e_type][1].append(d)
+        raw_subg_eset[e_type].append(edge_set)
+        raw_subg_etype[e_type].append(r)
+
+        if reverse is True:
+            r_type = str(r + num_rels)
+            re_type = (d_type, r_type, s_type)
+            if raw_reverse_sugb.get(re_type, None) is None:
+                raw_reverse_sugb[re_type] = ([], [])
+                raw_reverse_subg_etype[re_type] = []
+                raw_reverse_subg_eset[re_type] = []
+            raw_reverse_sugb[re_type][0].append(d)
+            raw_reverse_sugb[re_type][1].append(s)
+            raw_reverse_subg_eset[re_type].append(edge_set)
+            raw_reverse_subg_etype[re_type].append(r + num_rels)
+
+    for edge in train:
+        s, r, d = edge
+        assert r < num_rels
+        add_edge(s, r, d, reverse, 1)  # train set
+
+    for edge in valid:
+        s, r, d = edge
+        assert r < num_rels
+        add_edge(s, r, d, reverse, 2)  # valid set
+
+    for edge in test:
+        s, r, d = edge
+        assert r < num_rels
+        add_edge(s, r, d, reverse, 3)  # test set
+
+    subg = []
+    fg_s = []
+    fg_d = []
+    fg_etype = []
+    fg_settype = []
+    for e_type, val in raw_subg.items():
+        s, d = val
+        s = np.asarray(s)
+        d = np.asarray(d)
+        etype = raw_subg_etype[e_type]
+        etype = np.asarray(etype)
+        settype = raw_subg_eset[e_type]
+        settype = np.asarray(settype)
+
+        fg_s.append(s)
+        fg_d.append(d)
+        fg_etype.append(etype)
+        fg_settype.append(settype)
+
+    settype = np.concatenate(fg_settype)
+    if reverse is True:
+        settype = np.concatenate([settype, np.full((settype.shape[0]), 0)])
+    train_edge_mask = generate_mask_tensor(settype == 1)
+    valid_edge_mask = generate_mask_tensor(settype == 2)
+    test_edge_mask = generate_mask_tensor(settype == 3)
+
+    for e_type, val in raw_reverse_sugb.items():
+        s, d = val
+        s = np.asarray(s)
+        d = np.asarray(d)
+        etype = raw_reverse_subg_etype[e_type]
+        etype = np.asarray(etype)
+        settype = raw_reverse_subg_eset[e_type]
+        settype = np.asarray(settype)
+
+        fg_s.append(s)
+        fg_d.append(d)
+        fg_etype.append(etype)
+        fg_settype.append(settype)
+
+    s = np.concatenate(fg_s)
+    d = np.concatenate(fg_d)
+    g = dgl_graph((s, d), num_nodes=num_nodes)
+    etype = np.concatenate(fg_etype)
+    settype = np.concatenate(fg_settype)
+    etype = F.tensor(etype, dtype=F.data_type_dict["int64"])
+    train_edge_mask = train_edge_mask
+    valid_edge_mask = valid_edge_mask
+    test_edge_mask = test_edge_mask
+    train_mask = (
+        generate_mask_tensor(settype == 1)
+        if reverse is True
+        else train_edge_mask
+    )
+    valid_mask = (
+        generate_mask_tensor(settype == 2)
+        if reverse is True
+        else valid_edge_mask
+    )
+    test_mask = (
+        generate_mask_tensor(settype == 3)
+        if reverse is True
+        else test_edge_mask
+    )
+    ntype = F.full_1d(
+        num_nodes, 0, dtype=F.data_type_dict["int64"], ctx=F.cpu()
+    )
+
+    return g, (
+        etype,
+        ntype,
+        train_edge_mask,
+        valid_edge_mask,
+        test_edge_mask,
+        train_mask,
+        valid_mask,
+        test_mask,
+    )
+
+
+class FB15k237Dataset(KnowledgeGraphDataset):
+    r"""FB15k237 link prediction dataset.
+
+    FB15k-237 is a subset of FB15k where inverse
+    relations are removed. When creating the dataset,
+    a reverse edge with reversed relation types are
+    created for each edge by default.
+
+    FB15k237 dataset statistics:
+
+    - Nodes: 14541
+    - Number of relation types: 237
+    - Number of reversed relation types: 237
+    - Label Split:
+
+        - Train: 272115
+        - Valid: 17535
+        - Test: 20466
+
+    Parameters
+    ----------
+    reverse : bool
+        Whether to add reverse edge. Default True.
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    num_nodes: int
+        Number of nodes
+    num_rels: int
+        Number of relation types
+
+    Examples
+    ----------
+    >>> dataset = FB15k237Dataset()
+    >>> g = dataset.graph
+    >>> e_type = g.edata['e_type']
+    >>>
+    >>> # get data split
+    >>> train_mask = g.edata['train_mask']
+    >>> val_mask = g.edata['val_mask']
+    >>> test_mask = g.edata['test_mask']
+    >>>
+    >>> train_set = th.arange(g.num_edges())[train_mask]
+    >>> val_set = th.arange(g.num_edges())[val_mask]
+    >>>
+    >>> # build train_g
+    >>> train_edges = train_set
+    >>> train_g = g.edge_subgraph(train_edges,
+                                  relabel_nodes=False)
+    >>> train_g.edata['e_type'] = e_type[train_edges];
+    >>>
+    >>> # build val_g
+    >>> val_edges = th.cat([train_edges, val_edges])
+    >>> val_g = g.edge_subgraph(val_edges,
+                                relabel_nodes=False)
+    >>> val_g.edata['e_type'] = e_type[val_edges];
+    >>>
+    >>> # Train, Validation and Test
+    """
+
+    def __init__(
+        self,
+        reverse=True,
+        raw_dir=None,
+        force_reload=False,
+        verbose=True,
+        transform=None,
+    ):
+        name = "FB15k-237"
+        super(FB15k237Dataset, self).__init__(
+            name, reverse, raw_dir, force_reload, verbose, transform
+        )
+
+    def __getitem__(self, idx):
+        r"""Gets the graph object
+
+        Parameters
+        -----------
+        idx: int
+            Item index, FB15k237Dataset has only one graph object
+
+        Return
+        -------
+        :class:`dgl.DGLGraph`
+
+            The graph contains
+
+            - ``edata['e_type']``: edge relation type
+            - ``edata['train_edge_mask']``: positive training edge mask
+            - ``edata['val_edge_mask']``: positive validation edge mask
+            - ``edata['test_edge_mask']``: positive testing edge mask
+            - ``edata['train_mask']``: training edge set mask (include reversed training edges)
+            - ``edata['val_mask']``: validation edge set mask (include reversed validation edges)
+            - ``edata['test_mask']``: testing edge set mask (include reversed testing edges)
+            - ``ndata['ntype']``: node type. All 0 in this dataset
+        """
+        return super(FB15k237Dataset, self).__getitem__(idx)
+
+    def __len__(self):
+        r"""The number of graphs in the dataset."""
+        return super(FB15k237Dataset, self).__len__()
+
+
+class FB15kDataset(KnowledgeGraphDataset):
+    r"""FB15k link prediction dataset.
+
+    The FB15K dataset was introduced in `Translating Embeddings for Modeling
+    Multi-relational Data <http://papers.nips.cc/paper/5071-translating-embeddings-for-modeling-multi-relational-data.pdf>`_.
+    It is a subset of Freebase which contains about
+    14,951 entities with 1,345 different relations.
+    When creating the dataset, a reverse edge with
+    reversed relation types are created for each edge
+    by default.
+
+    FB15k dataset statistics:
+
+    - Nodes: 14,951
+    - Number of relation types: 1,345
+    - Number of reversed relation types: 1,345
+    - Label Split:
+
+        - Train: 483142
+        - Valid: 50000
+        - Test: 59071
+
+    Parameters
+    ----------
+    reverse : bool
+        Whether to add reverse edge. Default True.
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    num_nodes: int
+        Number of nodes
+    num_rels: int
+        Number of relation types
+
+    Examples
+    ----------
+    >>> dataset = FB15kDataset()
+    >>> g = dataset.graph
+    >>> e_type = g.edata['e_type']
+    >>>
+    >>> # get data split
+    >>> train_mask = g.edata['train_mask']
+    >>> val_mask = g.edata['val_mask']
+    >>>
+    >>> train_set = th.arange(g.num_edges())[train_mask]
+    >>> val_set = th.arange(g.num_edges())[val_mask]
+    >>>
+    >>> # build train_g
+    >>> train_edges = train_set
+    >>> train_g = g.edge_subgraph(train_edges,
+                                  relabel_nodes=False)
+    >>> train_g.edata['e_type'] = e_type[train_edges];
+    >>>
+    >>> # build val_g
+    >>> val_edges = th.cat([train_edges, val_edges])
+    >>> val_g = g.edge_subgraph(val_edges,
+                                relabel_nodes=False)
+    >>> val_g.edata['e_type'] = e_type[val_edges];
+    >>>
+    >>> # Train, Validation and Test
+    >>>
+    """
+
+    def __init__(
+        self,
+        reverse=True,
+        raw_dir=None,
+        force_reload=False,
+        verbose=True,
+        transform=None,
+    ):
+        name = "FB15k"
+        super(FB15kDataset, self).__init__(
+            name, reverse, raw_dir, force_reload, verbose, transform
+        )
+
+    def __getitem__(self, idx):
+        r"""Gets the graph object
+
+        Parameters
+        -----------
+        idx: int
+            Item index, FB15kDataset has only one graph object
+
+        Return
+        -------
+        :class:`dgl.DGLGraph`
+
+            The graph contains
+
+            - ``edata['e_type']``: edge relation type
+            - ``edata['train_edge_mask']``: positive training edge mask
+            - ``edata['val_edge_mask']``: positive validation edge mask
+            - ``edata['test_edge_mask']``: positive testing edge mask
+            - ``edata['train_mask']``: training edge set mask (include reversed training edges)
+            - ``edata['val_mask']``: validation edge set mask (include reversed validation edges)
+            - ``edata['test_mask']``: testing edge set mask (include reversed testing edges)
+            - ``ndata['ntype']``: node type. All 0 in this dataset
+        """
+        return super(FB15kDataset, self).__getitem__(idx)
+
+    def __len__(self):
+        r"""The number of graphs in the dataset."""
+        return super(FB15kDataset, self).__len__()
+
+
+class WN18Dataset(KnowledgeGraphDataset):
+    r"""WN18 link prediction dataset.
+
+    The WN18 dataset was introduced in `Translating Embeddings for Modeling
+    Multi-relational Data <http://papers.nips.cc/paper/5071-translating-embeddings-for-modeling-multi-relational-data.pdf>`_.
+    It included the full 18 relations scraped from
+    WordNet for roughly 41,000 synsets. When creating
+    the dataset, a reverse edge with reversed relation
+    types are created for each edge by default.
+
+    WN18 dataset statistics:
+
+    - Nodes: 40943
+    - Number of relation types: 18
+    - Number of reversed relation types: 18
+    - Label Split:
+
+        - Train: 141442
+        - Valid: 5000
+        - Test: 5000
+
+    Parameters
+    ----------
+    reverse : bool
+        Whether to add reverse edge. Default True.
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    num_nodes: int
+        Number of nodes
+    num_rels: int
+        Number of relation types
+
+    Examples
+    ----------
+    >>> dataset = WN18Dataset()
+    >>> g = dataset.graph
+    >>> e_type = g.edata['e_type']
+    >>>
+    >>> # get data split
+    >>> train_mask = g.edata['train_mask']
+    >>> val_mask = g.edata['val_mask']
+    >>>
+    >>> train_set = th.arange(g.num_edges())[train_mask]
+    >>> val_set = th.arange(g.num_edges())[val_mask]
+    >>>
+    >>> # build train_g
+    >>> train_edges = train_set
+    >>> train_g = g.edge_subgraph(train_edges,
+                                  relabel_nodes=False)
+    >>> train_g.edata['e_type'] = e_type[train_edges];
+    >>>
+    >>> # build val_g
+    >>> val_edges = th.cat([train_edges, val_edges])
+    >>> val_g = g.edge_subgraph(val_edges,
+                                relabel_nodes=False)
+    >>> val_g.edata['e_type'] = e_type[val_edges];
+    >>>
+    >>> # Train, Validation and Test
+    >>>
+    """
+
+    def __init__(
+        self,
+        reverse=True,
+        raw_dir=None,
+        force_reload=False,
+        verbose=True,
+        transform=None,
+    ):
+        name = "wn18"
+        super(WN18Dataset, self).__init__(
+            name, reverse, raw_dir, force_reload, verbose, transform
+        )
+
+    def __getitem__(self, idx):
+        r"""Gets the graph object
+
+        Parameters
+        -----------
+        idx: int
+            Item index, WN18Dataset has only one graph object
+
+        Return
+        -------
+        :class:`dgl.DGLGraph`
+
+            The graph contains
+
+            - ``edata['e_type']``: edge relation type
+            - ``edata['train_edge_mask']``: positive training edge mask
+            - ``edata['val_edge_mask']``: positive validation edge mask
+            - ``edata['test_edge_mask']``: positive testing edge mask
+            - ``edata['train_mask']``: training edge set mask (include reversed training edges)
+            - ``edata['val_mask']``: validation edge set mask (include reversed validation edges)
+            - ``edata['test_mask']``: testing edge set mask (include reversed testing edges)
+            - ``ndata['ntype']``: node type. All 0 in this dataset
+        """
+        return super(WN18Dataset, self).__getitem__(idx)
+
+    def __len__(self):
+        r"""The number of graphs in the dataset."""
+        return super(WN18Dataset, self).__len__()
+
+
+def load_data(dataset):
+    r"""Load knowledge graph dataset for RGCN link prediction tasks
+
+    It supports three datasets: wn18, FB15k and FB15k-237
+
+    Parameters
+    ----------
+    dataset: str
+        The name of the dataset to load.
+
+    Return
+    ------
+    The dataset object.
+    """
+    if dataset == "wn18":
+        return WN18Dataset()
+    elif dataset == "FB15k":
+        return FB15kDataset()
+    elif dataset == "FB15k-237":
+        return FB15k237Dataset()

python/dgl/data/minigc.py

@@ -0,0 +1,248 @@
+"""A mini synthetic dataset for graph classification benchmark."""
+import math
+import os
+
+import networkx as nx
+import numpy as np
+
+from .. import backend as F
+from ..convert import from_networkx
+from ..transforms import add_self_loop
+from .dgl_dataset import DGLDataset
+from .utils import load_graphs, makedirs, save_graphs
+
+__all__ = ["MiniGCDataset"]
+
+
+class MiniGCDataset(DGLDataset):
+    """The synthetic graph classification dataset class.
+
+    The datset contains 8 different types of graphs.
+
+    - class 0 : cycle graph
+    - class 1 : star graph
+    - class 2 : wheel graph
+    - class 3 : lollipop graph
+    - class 4 : hypercube graph
+    - class 5 : grid graph
+    - class 6 : clique graph
+    - class 7 : circular ladder graph
+
+    Parameters
+    ----------
+    num_graphs: int
+        Number of graphs in this dataset.
+    min_num_v: int
+        Minimum number of nodes for graphs
+    max_num_v: int
+        Maximum number of nodes for graphs
+    seed: int, default is 0
+        Random seed for data generation
+    transform: callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    num_graphs : int
+        Number of graphs
+    min_num_v : int
+        The minimum number of nodes
+    max_num_v : int
+        The maximum number of nodes
+    num_classes : int
+        The number of classes
+
+    Examples
+    --------
+    >>> data = MiniGCDataset(100, 16, 32, seed=0)
+
+    The dataset instance is an iterable
+
+    >>> len(data)
+    100
+    >>> g, label = data[64]
+    >>> g
+    Graph(num_nodes=20, num_edges=82,
+          ndata_schemes={}
+          edata_schemes={})
+    >>> label
+    tensor(5)
+
+    Batch the graphs and labels for mini-batch training
+
+    >>> graphs, labels = zip(*[data[i] for i in range(16)])
+    >>> batched_graphs = dgl.batch(graphs)
+    >>> batched_labels = torch.tensor(labels)
+    >>> batched_graphs
+    Graph(num_nodes=356, num_edges=1060,
+          ndata_schemes={}
+          edata_schemes={})
+    """
+
+    def __init__(
+        self,
+        num_graphs,
+        min_num_v,
+        max_num_v,
+        seed=0,
+        save_graph=True,
+        force_reload=False,
+        verbose=False,
+        transform=None,
+    ):
+        self.num_graphs = num_graphs
+        self.min_num_v = min_num_v
+        self.max_num_v = max_num_v
+        self.seed = seed
+        self.save_graph = save_graph
+
+        super(MiniGCDataset, self).__init__(
+            name="minigc",
+            hash_key=(num_graphs, min_num_v, max_num_v, seed),
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        self.graphs = []
+        self.labels = []
+        self._generate(self.seed)
+
+    def __len__(self):
+        """Return the number of graphs in the dataset."""
+        return len(self.graphs)
+
+    def __getitem__(self, idx):
+        """Get the idx-th sample.
+
+        Parameters
+        ---------
+        idx : int
+            The sample index.
+
+        Returns
+        -------
+        (:class:`dgl.Graph`, Tensor)
+            The graph and its label.
+        """
+        if self._transform is None:
+            g = self.graphs[idx]
+        else:
+            g = self._transform(self.graphs[idx])
+        return g, self.labels[idx]
+
+    def has_cache(self):
+        graph_path = os.path.join(
+            self.save_path, "dgl_graph_{}.bin".format(self.hash)
+        )
+        if os.path.exists(graph_path):
+            return True
+
+        return False
+
+    def save(self):
+        """save the graph list and the labels"""
+        if self.save_graph:
+            graph_path = os.path.join(
+                self.save_path, "dgl_graph_{}.bin".format(self.hash)
+            )
+            save_graphs(str(graph_path), self.graphs, {"labels": self.labels})
+
+    def load(self):
+        graphs, label_dict = load_graphs(
+            os.path.join(self.save_path, "dgl_graph_{}.bin".format(self.hash))
+        )
+        self.graphs = graphs
+        self.labels = label_dict["labels"]
+
+    @property
+    def num_classes(self):
+        """Number of classes."""
+        return 8
+
+    def _generate(self, seed):
+        if seed is not None:
+            np.random.seed(seed)
+        self._gen_cycle(self.num_graphs // 8)
+        self._gen_star(self.num_graphs // 8)
+        self._gen_wheel(self.num_graphs // 8)
+        self._gen_lollipop(self.num_graphs // 8)
+        self._gen_hypercube(self.num_graphs // 8)
+        self._gen_grid(self.num_graphs // 8)
+        self._gen_clique(self.num_graphs // 8)
+        self._gen_circular_ladder(self.num_graphs - len(self.graphs))
+        # preprocess
+        for i in range(self.num_graphs):
+            # convert to DGLGraph, and add self loops
+            self.graphs[i] = add_self_loop(from_networkx(self.graphs[i]))
+        self.labels = F.tensor(np.array(self.labels).astype(np.int64))
+
+    def _gen_cycle(self, n):
+        for _ in range(n):
+            num_v = np.random.randint(self.min_num_v, self.max_num_v)
+            g = nx.cycle_graph(num_v)
+            self.graphs.append(g)
+            self.labels.append(0)
+
+    def _gen_star(self, n):
+        for _ in range(n):
+            num_v = np.random.randint(self.min_num_v, self.max_num_v)
+            # nx.star_graph(N) gives a star graph with N+1 nodes
+            g = nx.star_graph(num_v - 1)
+            self.graphs.append(g)
+            self.labels.append(1)
+
+    def _gen_wheel(self, n):
+        for _ in range(n):
+            num_v = np.random.randint(self.min_num_v, self.max_num_v)
+            g = nx.wheel_graph(num_v)
+            self.graphs.append(g)
+            self.labels.append(2)
+
+    def _gen_lollipop(self, n):
+        for _ in range(n):
+            num_v = np.random.randint(self.min_num_v, self.max_num_v)
+            path_len = np.random.randint(2, num_v // 2)
+            g = nx.lollipop_graph(m=num_v - path_len, n=path_len)
+            self.graphs.append(g)
+            self.labels.append(3)
+
+    def _gen_hypercube(self, n):
+        for _ in range(n):
+            num_v = np.random.randint(self.min_num_v, self.max_num_v)
+            g = nx.hypercube_graph(int(math.log(num_v, 2)))
+            g = nx.convert_node_labels_to_integers(g)
+            self.graphs.append(g)
+            self.labels.append(4)
+
+    def _gen_grid(self, n):
+        for _ in range(n):
+            num_v = np.random.randint(self.min_num_v, self.max_num_v)
+            assert num_v >= 4, (
+                "We require a grid graph to contain at least two "
+                "rows and two columns, thus 4 nodes, got {:d} "
+                "nodes".format(num_v)
+            )
+            n_rows = np.random.randint(2, num_v // 2)
+            n_cols = num_v // n_rows
+            g = nx.grid_graph([n_rows, n_cols])
+            g = nx.convert_node_labels_to_integers(g)
+            self.graphs.append(g)
+            self.labels.append(5)
+
+    def _gen_clique(self, n):
+        for _ in range(n):
+            num_v = np.random.randint(self.min_num_v, self.max_num_v)
+            g = nx.complete_graph(num_v)
+            self.graphs.append(g)
+            self.labels.append(6)
+
+    def _gen_circular_ladder(self, n):
+        for _ in range(n):
+            num_v = np.random.randint(self.min_num_v, self.max_num_v)
+            g = nx.circular_ladder_graph(num_v // 2)
+            self.graphs.append(g)
+            self.labels.append(7)

python/dgl/data/movielens.py

@@ -0,0 +1,646 @@
+"""MovieLens dataset"""
+import os
+
+import numpy as np
+import pandas as pd
+
+from torch import LongTensor, Tensor
+
+from ..base import dgl_warning
+from ..convert import heterograph
+from .dgl_dataset import DGLDataset
+
+from .utils import (
+    _get_dgl_url,
+    download,
+    extract_archive,
+    load_graphs,
+    load_info,
+    save_graphs,
+    save_info,
+    split_dataset,
+)
+
+GENRES_ML_100K = [
+    "unknown",
+    "Action",
+    "Adventure",
+    "Animation",
+    "Children",
+    "Comedy",
+    "Crime",
+    "Documentary",
+    "Drama",
+    "Fantasy",
+    "Film-Noir",
+    "Horror",
+    "Musical",
+    "Mystery",
+    "Romance",
+    "Sci-Fi",
+    "Thriller",
+    "War",
+    "Western",
+]
+GENRES_ML_1M = GENRES_ML_100K[1:]
+GENRES_ML_10M = GENRES_ML_100K + ["IMAX"]
+
+try:
+    import torch
+except ImportError:
+    HAS_TORCH = False
+else:
+    HAS_TORCH = True
+
+
+def check_pytorch():
+    """Check if PyTorch is the backend."""
+    if not HAS_TORCH:
+        raise ModuleNotFoundError(
+            "MovieLensDataset requires PyTorch to be the backend."
+        )
+
+
+class MovieLensDataset(DGLDataset):
+    r"""MovieLens dataset for edge prediction tasks. The raw datasets are extracted from
+    `MovieLens <https://grouplens.org/datasets/movielens/>`, introduced by
+    `Movielens unplugged: experiences with an occasionally connected recommender system <https://dl.acm.org/doi/10.1145/604045.604094>`.
+
+    The datasets consist of user ratings for movies and incorporate additional user/movie information in the form of features.
+    The nodes represent users and movies, and the edges store ratings that users assign to movies.
+
+    Statistics:
+
+    MovieLens-100K (ml-100k)
+
+    - Users: 943
+    - Movies: 1,682
+    - Ratings: 100,000 (1, 2, 3, 4, 5)
+
+    MovieLens-1M (ml-1m)
+
+    - Users: 6,040
+    - Movies: 3,706
+    - Ratings: 1,000,209 (1, 2, 3, 4, 5)
+
+    MovieLens-10M (ml-10m)
+
+    - Users: 69,878
+    - Movies: 10,677
+    - Ratings: 10,000,054 (0.5, 1, 1.5, ..., 4.5, 5.0)
+
+    Parameters
+    ----------
+    name: str
+        Dataset name. (:obj:`"ml-100k"`, :obj:`"ml-1m"`, :obj:`"ml-10m"`).
+    valid_ratio: int
+        Ratio of validation samples out of the whole dataset. Should be in (0.0, 1.0).
+    test_ratio: int, optional
+        Ratio of testing samples out of the whole dataset. Should be in (0.0, 1.0). And its sum with
+        :obj:`valid_ratio` should be in (0.0, 1.0) as well. This parameter is invalid
+        when :obj:`name` is :obj:`"ml-100k"`, since its testing samples are pre-specified.
+        Default: None
+    raw_dir : str, optional
+        Raw file directory to download/store the data.
+        Default: ~/.dgl/
+    force_reload : bool, optional
+        Whether to re-download(if the dataset has not been downloaded) and re-process the dataset.
+        Default: False
+    verbose : bool, optional
+        Whether to print progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+    random_state : int, optional
+        Random seed used for random dataset split. Default: 0
+
+    Notes
+    -----
+    - When :obj:`name` is :obj:`"ml-100k"`, the :obj:`test_ratio` is invalid, and the training ratio is equal to 1-:obj:`valid_ratio`.
+    When :obj:`name` is :obj:`"ml-1m"` or :obj:`"ml-10m"`, the :obj:`test_ratio` is valid,
+    and the training ratio is equal to 1-:obj:`valid_ratio`-:obj:`test_ratio`.
+    - The number of edges is doubled to form an undirected(bidirected) graph structure.
+
+    Examples
+    --------
+    >>> from dgl.data import MovieLensDataset
+    >>> dataset = MovieLensDataset(name='ml-100k', valid_ratio=0.2)
+    >>> g = dataset[0]
+    >>> g
+    Graph(num_nodes={'movie': 1682, 'user': 943},
+          num_edges={('movie', 'movie-user', 'user'): 100000, ('user', 'user-movie', 'movie'): 100000},
+          metagraph=[('movie', 'user', 'movie-user'), ('user', 'movie', 'user-movie')])
+
+    >>> # get ratings of edges in the training graph.
+    >>> rate = g.edges['user-movie'].data['rate'] # or rate = g.edges['movie-user'].data['rate']
+    >>> rate
+    tensor([5., 5., 3.,  ..., 3., 3., 5.])
+
+    >>> # get train, valid and test mask of edges
+    >>> train_mask = g.edges['user-movie'].data['train_mask']
+    >>> valid_mask = g.edges['user-movie'].data['valid_mask']
+    >>> test_mask = g.edges['user-movie'].data['test_mask']
+
+    >>> # get train, valid and test ratings
+    >>> train_ratings = rate[train_mask]
+    >>> valid_ratings = rate[valid_mask]
+    >>> test_ratings = rate[test_mask]
+
+    >>> # get input features of users
+    >>> g.nodes["user"].data["feat"] # or g.nodes["movie"].data["feat"] for movie nodes
+    tensor([[0.4800, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],
+            [1.0600, 1.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],
+            [0.4600, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],
+            ...,
+            [0.4000, 0.0000, 1.0000,  ..., 0.0000, 0.0000, 0.0000],
+            [0.9600, 1.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],
+            [0.4400, 0.0000, 1.0000,  ..., 0.0000, 0.0000, 0.0000]])
+
+    """
+
+    _url = {
+        "ml-100k": "dataset/ml-100k.zip",
+        "ml-1m": "dataset/ml-1m.zip",
+        "ml-10m": "dataset/ml-10m.zip",
+    }
+
+    def __init__(
+        self,
+        name,
+        valid_ratio,
+        test_ratio=None,
+        raw_dir=None,
+        force_reload=None,
+        verbose=None,
+        transform=None,
+        random_state=0,
+    ):
+        check_pytorch()
+        assert name in [
+            "ml-100k",
+            "ml-1m",
+            "ml-10m",
+        ], f"currently movielens does not support {name}"
+
+        # test regarding valid and test split ratio
+        assert (
+            valid_ratio > 0.0 and valid_ratio < 1.0
+        ), f"valid_ratio {valid_ratio} must be in (0.0, 1.0)"
+
+        if name in ["ml-1m", "ml-10m"]:
+            assert (
+                test_ratio is not None and test_ratio > 0.0 and test_ratio < 1.0
+            ), f"test_ratio({test_ratio}) must be set to a value in (0.0, 1.0) when using ml-1m and ml-10m"
+            assert (
+                test_ratio + valid_ratio > 0.0
+                and test_ratio + valid_ratio < 1.0
+            ), f"test_ratio({test_ratio}) + valid_ratio({valid_ratio}) must be set to (0.0, 1.0) when using ml-1m and ml-10m"
+
+        if name == "ml-100k" and test_ratio is not None:
+            dgl_warning(
+                f"test_ratio ({test_ratio}) is not set to None for ml-100k. "
+                "Note that dataset split would not be affected by the test_ratio since "
+                "testing samples of ml-100k have been pre-specified."
+            )
+
+        self.valid_ratio = valid_ratio
+        self.test_ratio = test_ratio
+        self.random_state = random_state
+
+        if name == "ml-100k":
+            self.genres = GENRES_ML_100K
+        elif name == "ml-1m":
+            self.genres = GENRES_ML_1M
+        elif name == "ml-10m":
+            self.genres = GENRES_ML_10M
+        else:
+            raise NotImplementedError
+
+        super(MovieLensDataset, self).__init__(
+            name=name,
+            url=_get_dgl_url(self._url[name]),
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def check_version(self):
+        valid_ratio, test_ratio = load_info(self.version_path)
+        if self.valid_ratio == valid_ratio and (
+            self.test_ratio == test_ratio if self.name != "ml-100k" else True
+        ):
+            return True
+        else:
+            if self.name == "ml-100k":
+                print(
+                    f"The current valid ratio ({self.valid_ratio}) "
+                    "is not the same as the last setting "
+                    f"(valid: {valid_ratio}). "
+                    f"MovieLens {self.name} will be re-processed with the new dataset split setting."
+                )
+            else:
+                print(
+                    f"At least one of current valid ({self.valid_ratio}) and test ({self.test_ratio}) ratio "
+                    "are not the same as the last setting "
+                    f"(valid: {valid_ratio}, test: {test_ratio}). "
+                    f"MovieLens {self.name} will be re-processed with the new dataset split setting."
+                )
+            return False
+
+    def download(self):
+        zip_file_path = os.path.join(self.raw_dir, self.name + ".zip")
+        download(self.url, path=zip_file_path)
+        extract_archive(zip_file_path, self.raw_dir, overwrite=True)
+
+    def process(self):
+        print(f"Starting processing {self.name} ...")
+
+        # 0. loading movie features
+        movie_feat = load_info(
+            os.path.join(self.raw_path, "movie_feat.pkl")
+        ).to(torch.float)
+        # 1. dataset split: train + (valid + ) test
+        if self.name == "ml-100k":
+            train_rating_data = self._load_raw_rates(
+                os.path.join(self.raw_path, "u1.base"), "\t"
+            )
+            test_rating_data = self._load_raw_rates(
+                os.path.join(self.raw_path, "u1.test"), "\t"
+            )
+            indices = np.arange(len(train_rating_data))
+            train, valid, _ = split_dataset(
+                indices,
+                [1 - self.valid_ratio, self.valid_ratio, 0.0],
+                shuffle=True,
+                random_state=self.random_state,
+            )
+            train_rating_data, valid_rating_data = (
+                train_rating_data.iloc[train.indices],
+                train_rating_data.iloc[valid.indices],
+            )
+            all_rating_data = pd.concat(
+                [train_rating_data, valid_rating_data, test_rating_data]
+            )
+
+        elif self.name == "ml-1m" or self.name == "ml-10m":
+            all_rating_data = self._load_raw_rates(
+                os.path.join(self.raw_path, "ratings.dat"), "::"
+            )
+            indices = np.arange(len(all_rating_data))
+            train, valid, test = split_dataset(
+                indices,
+                [
+                    1 - self.valid_ratio - self.test_ratio,
+                    self.valid_ratio,
+                    self.test_ratio,
+                ],
+                shuffle=True,
+                random_state=self.random_state,
+            )
+            train_rating_data, valid_rating_data, test_rating_data = (
+                all_rating_data.iloc[train.indices],
+                all_rating_data.iloc[valid.indices],
+                all_rating_data.iloc[test.indices],
+            )
+
+        # 2. load user and movie data, and drop those unseen in rating_data
+        user_data = self._load_raw_user_data()
+        movie_data = self._load_raw_movie_data()
+        user_data = self._drop_unseen_nodes(
+            data_df=user_data,
+            col_name="id",
+            reserved_ids_set=set(all_rating_data["user_id"].values),
+        )
+        movie_data = self._drop_unseen_nodes(
+            data_df=movie_data,
+            col_name="id",
+            reserved_ids_set=set(all_rating_data["movie_id"].values),
+        )
+
+        user_feat = Tensor(self._process_user_feat(user_data))
+
+        # 3. generate rating pairs
+        # Map user/movie to the global id
+        self._global_user_id_map = {
+            ele: i for i, ele in enumerate(user_data["id"])
+        }
+        self._global_movie_id_map = {
+            ele: i for i, ele in enumerate(movie_data["id"])
+        }
+
+        # pair value is idx rather than id
+        u_indices, v_indices, labels = self._generate_pair_value(
+            all_rating_data
+        )
+        all_rating_pairs = (
+            LongTensor(u_indices),
+            LongTensor(v_indices),
+        )
+        all_rating_values = Tensor(labels)
+
+        graph = self.construct_g(
+            all_rating_pairs, all_rating_values, user_feat, movie_feat
+        )
+        self.graph = self.add_masks(
+            graph, train_rating_data, valid_rating_data, test_rating_data
+        )
+
+        print(f"End processing {self.name} ...")
+
+    def construct_g(self, rate_pairs, rate_values, user_feat, movie_feat):
+        g = heterograph(
+            {
+                ("user", "user-movie", "movie"): (rate_pairs[0], rate_pairs[1]),
+                ("movie", "movie-user", "user"): (rate_pairs[1], rate_pairs[0]),
+            }
+        )
+        ndata = {"user": user_feat, "movie": movie_feat}
+        edata = {"user-movie": rate_values, "movie-user": rate_values}
+        g.ndata["feat"] = ndata
+        g.edata["rate"] = edata
+        return g
+
+    def add_masks(
+        self, g, train_rating_data, valid_rating_data, test_rating_data
+    ):
+        train_u_indices, train_v_indices, _ = self._generate_pair_value(
+            train_rating_data
+        )
+        valid_u_indices, valid_v_indices, _ = self._generate_pair_value(
+            valid_rating_data
+        )
+        test_u_indices, test_v_indices, _ = self._generate_pair_value(
+            test_rating_data
+        )
+
+        # user-movie
+        train_mask = torch.zeros((g.num_edges("user-movie"),), dtype=torch.bool)
+        train_mask[
+            g.edge_ids(train_u_indices, train_v_indices, etype="user-movie")
+        ] = True
+        valid_mask = torch.zeros((g.num_edges("user-movie"),), dtype=torch.bool)
+        valid_mask[
+            g.edge_ids(valid_u_indices, valid_v_indices, etype="user-movie")
+        ] = True
+        test_mask = torch.zeros((g.num_edges("user-movie"),), dtype=torch.bool)
+        test_mask[
+            g.edge_ids(test_u_indices, test_v_indices, etype="user-movie")
+        ] = True
+
+        g.edges["user-movie"].data["train_mask"] = train_mask
+        g.edges["user-movie"].data["valid_mask"] = valid_mask
+        g.edges["user-movie"].data["test_mask"] = test_mask
+
+        # movie-user
+        train_mask_rev = torch.zeros(
+            (g.num_edges("movie-user"),), dtype=torch.bool
+        )
+        train_mask_rev[
+            g.edge_ids(train_v_indices, train_u_indices, etype="movie-user")
+        ] = True
+        valid_mask_rev = torch.zeros(
+            (g.num_edges("movie-user"),), dtype=torch.bool
+        )
+        valid_mask_rev[
+            g.edge_ids(valid_v_indices, valid_u_indices, etype="movie-user")
+        ] = True
+        test_mask_rev = torch.zeros(
+            (g.num_edges("movie-user"),), dtype=torch.bool
+        )
+        test_mask_rev[
+            g.edge_ids(test_v_indices, test_u_indices, etype="movie-user")
+        ] = True
+
+        g.edges["movie-user"].data["train_mask"] = train_mask_rev
+        g.edges["movie-user"].data["valid_mask"] = valid_mask_rev
+        g.edges["movie-user"].data["test_mask"] = test_mask_rev
+
+        return g
+
+    def has_cache(self):
+        if (
+            os.path.exists(self.graph_path)
+            and os.path.exists(self.version_path)
+            and self.check_version()
+        ):
+            return True
+        return False
+
+    def save(self):
+        save_graphs(self.graph_path, [self.graph])
+        save_info(self.version_path, [self.valid_ratio, self.test_ratio])
+        if self.verbose:
+            print(f"Done saving data into {self.raw_path}.")
+
+    def load(self):
+        g_list, _ = load_graphs(self.graph_path)
+        self.graph = g_list[0]
+
+        """
+        To avoid the problem each time loading boolean tensor from the disk, boolean values
+        would be automatically converted into torch.uint8 types, and a deprecation warning would
+        be raised for using torch.uint8
+        """
+        for e in self.graph.etypes:
+            self.graph.edges[e].data["train_mask"] = (
+                self.graph.edges[e].data["train_mask"].to(torch.bool)
+            )
+            self.graph.edges[e].data["valid_mask"] = (
+                self.graph.edges[e].data["valid_mask"].to(torch.bool)
+            )
+            self.graph.edges[e].data["test_mask"] = (
+                self.graph.edges[e].data["test_mask"].to(torch.bool)
+            )
+
+    def __getitem__(self, idx):
+        assert (
+            idx == 0
+        ), "This dataset has only one set of training, validation and testing graph"
+        if self._transform is None:
+            return self.graph
+        else:
+            return self._transform(self.graph)
+
+    def __len__(self):
+        return 1
+
+    @property
+    def raw_path(self):
+        return os.path.join(self.raw_dir, self.name)
+
+    @property
+    def graph_path(self):
+        return os.path.join(self.raw_path, self.name + ".bin")
+
+    @property
+    def version_path(self):
+        return os.path.join(self.raw_path, self.name + "_version.pkl")
+
+    def _process_user_feat(self, user_data):
+        if self.name == "ml-100k" or self.name == "ml-1m":
+            ages = user_data["age"].values.astype(np.float32)
+            gender = (user_data["gender"] == "F").values.astype(np.float32)
+            all_occupations = set(user_data["occupation"])
+            occupation_map = {ele: i for i, ele in enumerate(all_occupations)}
+            occupation_one_hot = np.zeros(
+                shape=(user_data.shape[0], len(all_occupations)),
+                dtype=np.float32,
+            )
+            occupation_one_hot[
+                np.arange(user_data.shape[0]),
+                np.array(
+                    [occupation_map[ele] for ele in user_data["occupation"]]
+                ),
+            ] = 1
+            user_features = np.concatenate(
+                [
+                    ages.reshape((user_data.shape[0], 1)) / 50.0,
+                    gender.reshape((user_data.shape[0], 1)),
+                    occupation_one_hot,
+                ],
+                axis=1,
+            )
+        elif self.name == "ml-10m":
+            user_features = np.zeros(
+                shape=(user_data.shape[0], 1), dtype=np.float32
+            )
+        else:
+            raise NotImplementedError
+        return user_features
+
+    def _load_raw_user_data(self):
+        if self.name == "ml-100k":
+            user_data = pd.read_csv(
+                os.path.join(self.raw_path, "u.user"),
+                sep="|",
+                header=None,
+                names=["id", "age", "gender", "occupation", "zip_code"],
+                engine="python",
+            )
+        elif self.name == "ml-1m":
+            user_data = pd.read_csv(
+                os.path.join(self.raw_path, "users.dat"),
+                sep="::",
+                header=None,
+                names=["id", "gender", "age", "occupation", "zip_code"],
+                engine="python",
+            )
+        elif self.name == "ml-10m":
+            rating_info = pd.read_csv(
+                os.path.join(self.raw_path, "ratings.dat"),
+                sep="::",
+                header=None,
+                names=["user_id", "movie_id", "rating", "timestamp"],
+                dtype={
+                    "user_id": np.int32,
+                    "movie_id": np.int32,
+                    "ratings": np.float32,
+                    "timestamp": np.int64,
+                },
+                engine="python",
+            )
+            user_data = pd.DataFrame(
+                np.unique(rating_info["user_id"].values.astype(np.int32)),
+                columns=["id"],
+            )
+        else:
+            raise NotImplementedError
+        return user_data
+
+    def _load_raw_movie_data(self):
+        file_path = os.path.join(self.raw_path, "u.item")
+        if self.name == "ml-100k":
+            movie_data = pd.read_csv(
+                file_path,
+                sep="|",
+                header=None,
+                names=[
+                    "id",
+                    "title",
+                    "release_date",
+                    "video_release_date",
+                    "url",
+                ]
+                + GENRES_ML_100K,
+                engine="python",
+                encoding="ISO-8859-1",
+            )
+        elif self.name == "ml-1m" or self.name == "ml-10m":
+            file_path = os.path.join(self.raw_path, "movies.dat")
+            movie_data = pd.read_csv(
+                file_path,
+                sep="::",
+                header=None,
+                names=["id", "title", "genres"],
+                encoding="iso-8859-1",
+                engine="python",
+            )
+            genre_map = {ele: i for i, ele in enumerate(self.genres)}
+            genre_map["Children's"] = genre_map["Children"]
+            genre_map["Childrens"] = genre_map["Children"]
+            movie_genres = np.zeros(
+                shape=(movie_data.shape[0], len(self.genres)), dtype=np.float32
+            )
+            for i, genres in enumerate(movie_data["genres"]):
+                for ele in genres.split("|"):
+                    if ele in genre_map:
+                        movie_genres[i, genre_map[ele]] = 1.0
+                    else:
+                        movie_genres[i, genre_map["unknown"]] = 1.0
+            for idx, genre_name in enumerate(self.genres):
+                movie_data[genre_name] = movie_genres[:, idx]
+            movie_data = movie_data.drop(columns=["genres"])
+        else:
+            raise NotImplementedError
+
+        return movie_data
+
+    def _load_raw_rates(self, file_path, sep):
+        rating_data = pd.read_csv(
+            file_path,
+            sep=sep,
+            header=None,
+            names=["user_id", "movie_id", "rating", "timestamp"],
+            dtype={
+                "user_id": np.int32,
+                "movie_id": np.int32,
+                "ratings": np.float32,
+                "timestamp": np.int64,
+            },
+            engine="python",
+        )
+        rating_data = rating_data.reset_index(drop=True)
+        return rating_data
+
+    def _drop_unseen_nodes(self, data_df, col_name, reserved_ids_set):
+        data_df = data_df[data_df[col_name].isin(reserved_ids_set)]
+        data_df.reset_index(drop=True, inplace=True)
+        return data_df
+
+    def _generate_pair_value(self, rating_data):
+        rating_pairs = (
+            np.array(
+                [
+                    self._global_user_id_map[ele]
+                    for ele in rating_data["user_id"]
+                ],
+                dtype=np.int32,
+            ),
+            np.array(
+                [
+                    self._global_movie_id_map[ele]
+                    for ele in rating_data["movie_id"]
+                ],
+                dtype=np.int32,
+            ),
+        )
+        rating_values = rating_data["rating"].values.astype(np.float32)
+        return rating_pairs[0], rating_pairs[1], rating_values
+
+    def __repr__(self):
+        return (
+            f'Dataset("{self.name}", num_graphs={len(self)},'
+            + f" save_path={self.raw_path}), valid_ratio={self.valid_ratio}, test_ratio={self.test_ratio}"
+        )

python/dgl/data/pattern.py

@@ -0,0 +1,130 @@
+""" PATTERNDataset for inductive learning. """
+import os
+
+from .dgl_dataset import DGLBuiltinDataset
+from .utils import _get_dgl_url, load_graphs
+
+
+class PATTERNDataset(DGLBuiltinDataset):
+    r"""PATTERN dataset for graph pattern recognition task.
+
+    Each graph G contains 5 communities with sizes randomly selected between [5, 35].
+    The SBM of each community is p = 0.5, q = 0.35, and the node features on G are
+    generated with a uniform random distribution with a vocabulary of size 3, i.e. {0, 1, 2}.
+    Then randomly generate 100 patterns P composed of 20 nodes with intra-probability :math:`p_P` = 0.5
+    and extra-probability :math:`q_P` = 0.5 (i.e. 50% of nodes in P are connected to G). The node features
+    for P are also generated as a random signal with values {0, 1, 2}. The graphs are of sizes
+    44-188 nodes. The output node labels have value 1 if the node belongs to P and value 0 if it is in G.
+
+    Reference `<https://arxiv.org/pdf/2003.00982.pdf>`_
+
+    Statistics:
+
+    - Train examples: 10,000
+    - Valid examples: 2,000
+    - Test examples: 2,000
+    - Number of classes for each node: 2
+
+    Parameters
+    ----------
+    mode : str
+        Must be one of ('train', 'valid', 'test').
+        Default: 'train'
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset.
+        Default: False
+    verbose : bool
+        Whether to print out progress information.
+        Default: False
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of classes for each node.
+
+    Examples
+    --------
+    >>> from dgl.data import PATTERNDataset
+    >>> data = PATTERNDataset(mode='train')
+    >>> data.num_classes
+    2
+    >>> len(trainset)
+    10000
+    >>> data[0]
+    Graph(num_nodes=108, num_edges=4884, ndata_schemes={'feat': Scheme(shape=(), dtype=torch.int64), 'label': Scheme(shape=(), dtype=torch.int16)}
+    edata_schemes={'feat': Scheme(shape=(1,), dtype=torch.float32)})
+    """
+
+    def __init__(
+        self,
+        mode="train",
+        raw_dir=None,
+        force_reload=False,
+        verbose=False,
+        transform=None,
+    ):
+        assert mode in ["train", "valid", "test"]
+        self.mode = mode
+        _url = _get_dgl_url("dataset/SBM_PATTERN.zip")
+
+        super(PATTERNDataset, self).__init__(
+            name="pattern",
+            url=_url,
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        self.load()
+
+    @property
+    def graph_path(self):
+        return os.path.join(
+            self.save_path, "SBM_PATTERN_{}.bin".format(self.mode)
+        )
+
+    def has_cache(self):
+        return os.path.exists(self.graph_path)
+
+    def load(self):
+        self._graphs, _ = load_graphs(self.graph_path)
+
+    @property
+    def num_classes(self):
+        r"""Number of classes for each node."""
+        return 2
+
+    def __len__(self):
+        r"""The number of examples in the dataset."""
+        return len(self._graphs)
+
+    def __getitem__(self, idx):
+        r"""Get the idx^th sample.
+
+        Parameters
+        ---------
+        idx : int
+            The sample index.
+
+        Returns
+        -------
+        :class:`dgl.DGLGraph`
+            graph structure, node features, node labels and edge features.
+
+            - ``ndata['feat']``: node features
+            - ``ndata['label']``: node labels
+            - ``edata['feat']``: edge features
+        """
+        if self._transform is None:
+            return self._graphs[idx]
+        else:
+            return self._transform(self._graphs[idx])

python/dgl/data/ppi.py

@@ -0,0 +1,226 @@
+""" PPIDataset for inductive learning. """
+import json
+import os
+
+import networkx as nx
+import numpy as np
+from networkx.readwrite import json_graph
+
+from .. import backend as F
+from ..convert import from_networkx
+from .dgl_dataset import DGLBuiltinDataset
+from .utils import _get_dgl_url, load_graphs, load_info, save_graphs, save_info
+
+
+class PPIDataset(DGLBuiltinDataset):
+    r"""Protein-Protein Interaction dataset for inductive node classification
+
+    A toy Protein-Protein Interaction network dataset. The dataset contains
+    24 graphs. The average number of nodes per graph is 2372. Each node has
+    50 features and 121 labels. 20 graphs for training, 2 for validation
+    and 2 for testing.
+
+    Reference: `<http://snap.stanford.edu/graphsage/>`_
+
+    Statistics:
+
+    - Train examples: 20
+    - Valid examples: 2
+    - Test examples: 2
+
+    Parameters
+    ----------
+    mode : str
+        Must be one of ('train', 'valid', 'test').
+        Default: 'train'
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset.
+        Default: False
+    verbose : bool
+        Whether to print out progress information.
+        Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    num_labels : int
+        Number of labels for each node
+    labels : Tensor
+        Node labels
+    features : Tensor
+        Node features
+
+    Examples
+    --------
+    >>> dataset = PPIDataset(mode='valid')
+    >>> num_classes = dataset.num_classes
+    >>> for g in dataset:
+    ....    feat = g.ndata['feat']
+    ....    label = g.ndata['label']
+    ....    # your code here
+    >>>
+    """
+
+    def __init__(
+        self,
+        mode="train",
+        raw_dir=None,
+        force_reload=False,
+        verbose=False,
+        transform=None,
+    ):
+        assert mode in ["train", "valid", "test"]
+        self.mode = mode
+        _url = _get_dgl_url("dataset/ppi.zip")
+        super(PPIDataset, self).__init__(
+            name="ppi",
+            url=_url,
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        graph_file = os.path.join(
+            self.save_path, "{}_graph.json".format(self.mode)
+        )
+        label_file = os.path.join(
+            self.save_path, "{}_labels.npy".format(self.mode)
+        )
+        feat_file = os.path.join(
+            self.save_path, "{}_feats.npy".format(self.mode)
+        )
+        graph_id_file = os.path.join(
+            self.save_path, "{}_graph_id.npy".format(self.mode)
+        )
+
+        g_data = json.load(open(graph_file))
+        self._labels = np.load(label_file)
+        self._feats = np.load(feat_file)
+        self.graph = from_networkx(
+            nx.DiGraph(json_graph.node_link_graph(g_data))
+        )
+        graph_id = np.load(graph_id_file)
+
+        # lo, hi means the range of graph ids for different portion of the dataset,
+        # 20 graphs for training, 2 for validation and 2 for testing.
+        lo, hi = 1, 21
+        if self.mode == "valid":
+            lo, hi = 21, 23
+        elif self.mode == "test":
+            lo, hi = 23, 25
+
+        graph_masks = []
+        self.graphs = []
+        for g_id in range(lo, hi):
+            g_mask = np.where(graph_id == g_id)[0]
+            graph_masks.append(g_mask)
+            g = self.graph.subgraph(g_mask)
+            g.ndata["feat"] = F.tensor(
+                self._feats[g_mask], dtype=F.data_type_dict["float32"]
+            )
+            g.ndata["label"] = F.tensor(
+                self._labels[g_mask], dtype=F.data_type_dict["float32"]
+            )
+            self.graphs.append(g)
+
+    @property
+    def graph_list_path(self):
+        return os.path.join(
+            self.save_path, "{}_dgl_graph_list.bin".format(self.mode)
+        )
+
+    @property
+    def g_path(self):
+        return os.path.join(
+            self.save_path, "{}_dgl_graph.bin".format(self.mode)
+        )
+
+    @property
+    def info_path(self):
+        return os.path.join(self.save_path, "{}_info.pkl".format(self.mode))
+
+    def has_cache(self):
+        return (
+            os.path.exists(self.graph_list_path)
+            and os.path.exists(self.g_path)
+            and os.path.exists(self.info_path)
+        )
+
+    def save(self):
+        save_graphs(self.graph_list_path, self.graphs)
+        save_graphs(self.g_path, self.graph)
+        save_info(
+            self.info_path, {"labels": self._labels, "feats": self._feats}
+        )
+
+    def load(self):
+        self.graphs = load_graphs(self.graph_list_path)[0]
+        g, _ = load_graphs(self.g_path)
+        self.graph = g[0]
+        info = load_info(self.info_path)
+        self._labels = info["labels"]
+        self._feats = info["feats"]
+
+    @property
+    def num_labels(self):
+        return 121
+
+    @property
+    def num_classes(self):
+        return 121
+
+    def __len__(self):
+        """Return number of samples in this dataset."""
+        return len(self.graphs)
+
+    def __getitem__(self, item):
+        """Get the item^th sample.
+
+        Parameters
+        ---------
+        item : int
+            The sample index.
+
+        Returns
+        -------
+        :class:`dgl.DGLGraph`
+            graph structure, node features and node labels.
+
+            - ``ndata['feat']``: node features
+            - ``ndata['label']``: node labels
+        """
+        if self._transform is None:
+            return self.graphs[item]
+        else:
+            return self._transform(self.graphs[item])
+
+
+class LegacyPPIDataset(PPIDataset):
+    """Legacy version of PPI Dataset"""
+
+    def __getitem__(self, item):
+        """Get the item^th sample.
+
+        Paramters
+        ---------
+        idx : int
+            The sample index.
+
+        Returns
+        -------
+        (dgl.DGLGraph, Tensor, Tensor)
+            The graph, features and its label.
+        """
+        if self._transform is None:
+            g = self.graphs[item]
+        else:
+            g = self._transform(self.graphs[item])
+        return g, g.ndata["feat"], g.ndata["label"]

python/dgl/data/qm7b.py

@@ -0,0 +1,177 @@
+"""QM7b dataset for graph property prediction (regression)."""
+import os
+
+from scipy import io
+
+from .. import backend as F
+from ..convert import graph as dgl_graph
+
+from .dgl_dataset import DGLDataset
+from .utils import check_sha1, download, load_graphs, save_graphs
+
+
+class QM7bDataset(DGLDataset):
+    r"""QM7b dataset for graph property prediction (regression)
+
+    This dataset consists of 7,211 molecules with 14 regression targets.
+    Nodes means atoms and edges means bonds. Edge data 'h' means
+    the entry of Coulomb matrix.
+
+    Reference: `<http://quantum-machine.org/datasets/>`_
+
+    Statistics:
+
+    - Number of graphs: 7,211
+    - Number of regression targets: 14
+    - Average number of nodes: 15
+    - Average number of edges: 245
+    - Edge feature size: 1
+
+    Parameters
+    ----------
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    num_tasks : int
+        Number of prediction tasks
+    num_labels : int
+        (DEPRECATED, use num_tasks instead) Number of prediction tasks
+
+    Raises
+    ------
+    UserWarning
+        If the raw data is changed in the remote server by the author.
+
+    Examples
+    --------
+    >>> data = QM7bDataset()
+    >>> data.num_tasks
+    14
+    >>>
+    >>> # iterate over the dataset
+    >>> for g, label in data:
+    ...     edge_feat = g.edata['h']  # get edge feature
+    ...     # your code here...
+    ...
+    >>>
+    """
+
+    _url = (
+        "http://deepchem.io.s3-website-us-west-1.amazonaws.com/"
+        "datasets/qm7b.mat"
+    )
+    _sha1_str = "4102c744bb9d6fd7b40ac67a300e49cd87e28392"
+
+    def __init__(
+        self, raw_dir=None, force_reload=False, verbose=False, transform=None
+    ):
+        super(QM7bDataset, self).__init__(
+            name="qm7b",
+            url=self._url,
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        mat_path = os.path.join(self.raw_dir, self.name + ".mat")
+        self.graphs, self.label = self._load_graph(mat_path)
+
+    def _load_graph(self, filename):
+        data = io.loadmat(filename)
+        labels = F.tensor(data["T"], dtype=F.data_type_dict["float32"])
+        feats = data["X"]
+        num_graphs = labels.shape[0]
+        graphs = []
+        for i in range(num_graphs):
+            edge_list = feats[i].nonzero()
+            g = dgl_graph(edge_list)
+            g.edata["h"] = F.tensor(
+                feats[i][edge_list[0], edge_list[1]].reshape(-1, 1),
+                dtype=F.data_type_dict["float32"],
+            )
+            graphs.append(g)
+        return graphs, labels
+
+    def save(self):
+        """save the graph list and the labels"""
+        graph_path = os.path.join(self.save_path, "dgl_graph.bin")
+        save_graphs(str(graph_path), self.graphs, {"labels": self.label})
+
+    def has_cache(self):
+        graph_path = os.path.join(self.save_path, "dgl_graph.bin")
+        return os.path.exists(graph_path)
+
+    def load(self):
+        graphs, label_dict = load_graphs(
+            os.path.join(self.save_path, "dgl_graph.bin")
+        )
+        self.graphs = graphs
+        self.label = label_dict["labels"]
+
+    def download(self):
+        file_path = os.path.join(self.raw_dir, self.name + ".mat")
+        download(self.url, path=file_path)
+        if not check_sha1(file_path, self._sha1_str):
+            raise UserWarning(
+                "File {} is downloaded but the content hash does not match."
+                "The repo may be outdated or download may be incomplete. "
+                "Otherwise you can create an issue for it.".format(self.name)
+            )
+
+    @property
+    def num_tasks(self):
+        """Number of prediction tasks."""
+        return self.num_labels
+
+    @property
+    def num_labels(self):
+        """Number of prediction tasks."""
+        return 14
+
+    @property
+    def num_classes(self):
+        """Number of prediction tasks."""
+        return 14
+
+    def __getitem__(self, idx):
+        r"""Get graph and label by index
+
+        Parameters
+        ----------
+        idx : int
+            Item index
+
+        Returns
+        -------
+        (:class:`dgl.DGLGraph`, Tensor)
+        """
+        if self._transform is None:
+            g = self.graphs[idx]
+        else:
+            g = self._transform(self.graphs[idx])
+        return g, self.label[idx]
+
+    def __len__(self):
+        r"""Number of graphs in the dataset.
+
+        Return
+        -------
+        int
+        """
+        return len(self.graphs)
+
+
+QM7b = QM7bDataset

python/dgl/data/qm9.py

@@ -0,0 +1,231 @@
+"""QM9 dataset for graph property prediction (regression)."""
+import os
+
+import numpy as np
+import scipy.sparse as sp
+
+from .. import backend as F
+from ..convert import graph as dgl_graph
+from ..transforms import to_bidirected
+
+from .dgl_dataset import DGLDataset
+from .utils import _get_dgl_url, download
+
+
+class QM9Dataset(DGLDataset):
+    r"""QM9 dataset for graph property prediction (regression)
+
+    This dataset consists of 130,831 molecules with 12 regression targets.
+    Nodes correspond to atoms and edges correspond to close atom pairs.
+
+    This dataset differs from :class:`~dgl.data.QM9EdgeDataset` in the following aspects:
+        1. Edges in this dataset are purely distance-based.
+        2. It only provides atoms' coordinates and atomic numbers as node features
+        3. It only provides 12 regression targets.
+
+    Reference:
+
+    - `"Quantum-Machine.org" <http://quantum-machine.org/datasets/>`_,
+    - `"Directional Message Passing for Molecular Graphs" <https://arxiv.org/abs/2003.03123>`_
+
+    Statistics:
+
+    - Number of graphs: 130,831
+    - Number of regression targets: 12
+
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | Keys   | Property                         | Description                                                                       | Unit                                        |
+    +========+==================================+===================================================================================+=============================================+
+    | mu     | :math:`\mu`                      | Dipole moment                                                                     | :math:`\textrm{D}`                          |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | alpha  | :math:`\alpha`                   | Isotropic polarizability                                                          | :math:`{a_0}^3`                             |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | homo   | :math:`\epsilon_{\textrm{HOMO}}` | Highest occupied molecular orbital energy                                         | :math:`\textrm{eV}`                         |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | lumo   | :math:`\epsilon_{\textrm{LUMO}}` | Lowest unoccupied molecular orbital energy                                        | :math:`\textrm{eV}`                         |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | gap    | :math:`\Delta \epsilon`          | Gap between :math:`\epsilon_{\textrm{HOMO}}` and :math:`\epsilon_{\textrm{LUMO}}` | :math:`\textrm{eV}`                         |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | r2     | :math:`\langle R^2 \rangle`      | Electronic spatial extent                                                         | :math:`{a_0}^2`                             |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | zpve   | :math:`\textrm{ZPVE}`            | Zero point vibrational energy                                                     | :math:`\textrm{eV}`                         |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | U0     | :math:`U_0`                      | Internal energy at 0K                                                             | :math:`\textrm{eV}`                         |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | U      | :math:`U`                        | Internal energy at 298.15K                                                        | :math:`\textrm{eV}`                         |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | H      | :math:`H`                        | Enthalpy at 298.15K                                                               | :math:`\textrm{eV}`                         |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | G      | :math:`G`                        | Free energy at 298.15K                                                            | :math:`\textrm{eV}`                         |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | Cv     | :math:`c_{\textrm{v}}`           | Heat capavity at 298.15K                                                          | :math:`\frac{\textrm{cal}}{\textrm{mol K}}` |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+
+    Parameters
+    ----------
+    label_keys : list
+        Names of the regression property, which should be a subset of the keys in the table above.
+    cutoff : float
+        Cutoff distance for interatomic interactions, i.e. two atoms are connected in the corresponding graph if the distance between them is no larger than this.
+        Default: 5.0 Angstrom
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    num_tasks : int
+        Number of prediction tasks
+    num_labels : int
+        (DEPRECATED, use num_tasks instead) Number of prediction tasks
+
+    Raises
+    ------
+    UserWarning
+        If the raw data is changed in the remote server by the author.
+
+    Examples
+    --------
+    >>> data = QM9Dataset(label_keys=['mu', 'gap'], cutoff=5.0)
+    >>> data.num_tasks
+    2
+    >>>
+    >>> # iterate over the dataset
+    >>> for g, label in data:
+    ...     R = g.ndata['R'] # get coordinates of each atom
+    ...     Z = g.ndata['Z'] # get atomic numbers of each atom
+    ...     # your code here...
+    >>>
+    """
+
+    def __init__(
+        self,
+        label_keys,
+        cutoff=5.0,
+        raw_dir=None,
+        force_reload=False,
+        verbose=False,
+        transform=None,
+    ):
+        self.cutoff = cutoff
+        self.label_keys = label_keys
+        self._url = _get_dgl_url("dataset/qm9_eV.npz")
+
+        super(QM9Dataset, self).__init__(
+            name="qm9",
+            url=self._url,
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        npz_path = f"{self.raw_dir}/qm9_eV.npz"
+        data_dict = np.load(npz_path, allow_pickle=True)
+        # data_dict['N'] contains the number of atoms in each molecule.
+        # Atomic properties (Z and R) of all molecules are concatenated as single tensors,
+        # so you need this value to select the correct atoms for each molecule.
+        self.N = data_dict["N"]
+        self.R = data_dict["R"]
+        self.Z = data_dict["Z"]
+        self.label = np.stack(
+            [data_dict[key] for key in self.label_keys], axis=1
+        )
+        self.N_cumsum = np.concatenate([[0], np.cumsum(self.N)])
+
+    def download(self):
+        file_path = f"{self.raw_dir}/qm9_eV.npz"
+        if not os.path.exists(file_path):
+            download(self._url, path=file_path)
+
+    @property
+    def num_labels(self):
+        r"""
+        Returns
+        --------
+        int
+            Number of prediction tasks.
+        """
+        return self.label.shape[1]
+
+    @property
+    def num_classes(self):
+        r"""
+        Returns
+        --------
+        int
+            Number of prediction tasks.
+        """
+        return self.label.shape[1]
+
+    @property
+    def num_tasks(self):
+        r"""
+        Returns
+        --------
+        int
+            Number of prediction tasks.
+        """
+        return self.label.shape[1]
+
+    def __getitem__(self, idx):
+        r"""Get graph and label by index
+
+        Parameters
+        ----------
+        idx : int
+            Item index
+
+        Returns
+        -------
+        dgl.DGLGraph
+            The graph contains:
+
+            - ``ndata['R']``: the coordinates of each atom
+            - ``ndata['Z']``: the atomic number
+
+        Tensor
+            Property values of molecular graphs
+        """
+        label = F.tensor(self.label[idx], dtype=F.data_type_dict["float32"])
+        n_atoms = self.N[idx]
+        R = self.R[self.N_cumsum[idx] : self.N_cumsum[idx + 1]]
+        dist = np.linalg.norm(R[:, None, :] - R[None, :, :], axis=-1)
+        adj = sp.csr_matrix(dist <= self.cutoff) - sp.eye(
+            n_atoms, dtype=np.bool_
+        )
+        adj = adj.tocoo()
+        u, v = F.tensor(adj.row), F.tensor(adj.col)
+        g = dgl_graph((u, v))
+        g = to_bidirected(g)
+        g.ndata["R"] = F.tensor(R, dtype=F.data_type_dict["float32"])
+        g.ndata["Z"] = F.tensor(
+            self.Z[self.N_cumsum[idx] : self.N_cumsum[idx + 1]],
+            dtype=F.data_type_dict["int64"],
+        )
+
+        if self._transform is not None:
+            g = self._transform(g)
+
+        return g, label
+
+    def __len__(self):
+        r"""Number of graphs in the dataset.
+
+        Return
+        -------
+        int
+        """
+        return self.label.shape[0]
+
+
+QM9 = QM9Dataset

python/dgl/data/qm9_edge.py

@@ -0,0 +1,296 @@
+""" QM9 dataset for graph property prediction (regression) """
+
+import os
+
+import numpy as np
+
+from .. import backend as F
+from ..convert import graph as dgl_graph
+
+from .dgl_dataset import DGLDataset
+from .utils import _get_dgl_url, download, extract_archive
+
+
+class QM9EdgeDataset(DGLDataset):
+    r"""QM9Edge dataset for graph property prediction (regression)
+
+    This dataset consists of 130,831 molecules with 19 regression targets.
+    Nodes correspond to atoms and edges correspond to bonds.
+
+    This dataset differs from :class:`~dgl.data.QM9Dataset` in the following aspects:
+        1. It includes the bonds in a molecule in the edges of the corresponding graph while the edges in :class:`~dgl.data.QM9Dataset` are purely distance-based.
+        2. It provides edge features, and node features in addition to the atoms' coordinates and atomic numbers.
+        3. It provides another 7 regression tasks(from 12 to 19).
+
+    This class is built based on a preprocessed version of the dataset, and we provide the preprocessing datails `here <https://gist.github.com/hengruizhang98/a2da30213b2356fff18b25385c9d3cd2>`_.
+
+    Reference:
+
+    - `"MoleculeNet: A Benchmark for Molecular Machine Learning" <https://arxiv.org/abs/1703.00564>`_
+    - `"Neural Message Passing for Quantum Chemistry" <https://arxiv.org/abs/1704.01212>`_
+
+    For
+    Statistics:
+
+    - Number of graphs: 130,831.
+    - Number of regression targets: 19.
+
+    Node attributes:
+
+    - pos: the 3D coordinates of each atom.
+    - attr: the 11D atom features.
+
+    Edge attributes:
+
+    - edge_attr: the 4D bond features.
+
+    Regression targets:
+
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | Keys   | Property                         | Description                                                                       | Unit                                        |
+    +========+==================================+===================================================================================+=============================================+
+    | mu     | :math:`\mu`                      | Dipole moment                                                                     | :math:`\textrm{D}`                          |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | alpha  | :math:`\alpha`                   | Isotropic polarizability                                                          | :math:`{a_0}^3`                             |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | homo   | :math:`\epsilon_{\textrm{HOMO}}` | Highest occupied molecular orbital energy                                         | :math:`\textrm{eV}`                         |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | lumo   | :math:`\epsilon_{\textrm{LUMO}}` | Lowest unoccupied molecular orbital energy                                        | :math:`\textrm{eV}`                         |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | gap    | :math:`\Delta \epsilon`          | Gap between :math:`\epsilon_{\textrm{HOMO}}` and :math:`\epsilon_{\textrm{LUMO}}` | :math:`\textrm{eV}`                         |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | r2     | :math:`\langle R^2 \rangle`      | Electronic spatial extent                                                         | :math:`{a_0}^2`                             |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | zpve   | :math:`\textrm{ZPVE}`            | Zero point vibrational energy                                                     | :math:`\textrm{eV}`                         |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | U0     | :math:`U_0`                      | Internal energy at 0K                                                             | :math:`\textrm{eV}`                         |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | U      | :math:`U`                        | Internal energy at 298.15K                                                        | :math:`\textrm{eV}`                         |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | H      | :math:`H`                        | Enthalpy at 298.15K                                                               | :math:`\textrm{eV}`                         |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | G      | :math:`G`                        | Free energy at 298.15K                                                            | :math:`\textrm{eV}`                         |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | Cv     | :math:`c_{\textrm{v}}`           | Heat capavity at 298.15K                                                          | :math:`\frac{\textrm{cal}}{\textrm{mol K}}` |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | U0_atom| :math:`U_0^{\textrm{ATOM}}`      | Atomization energy at 0K                                                          | :math:`\textrm{eV}`                         |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | U_atom | :math:`U^{\textrm{ATOM}}`        | Atomization energy at 298.15K                                                     | :math:`\textrm{eV}`                         |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | H_atom | :math:`H^{\textrm{ATOM}}`        | Atomization enthalpy at 298.15K                                                   | :math:`\textrm{eV}`                         |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | G_atom | :math:`G^{\textrm{ATOM}}`        | Atomization free energy at 298.15K                                                | :math:`\textrm{eV}`                         |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | A      | :math:`A`                        | Rotational constant                                                               | :math:`\textrm{GHz}`                        |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | B      | :math:`B`                        | Rotational constant                                                               | :math:`\textrm{GHz}`                        |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | C      | :math:`C`                        | Rotational constant                                                               | :math:`\textrm{GHz}`                        |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+
+    Parameters
+    ----------
+    label_keys : list
+        Names of the regression property, which should be a subset of the keys in the table above.
+        If not provided, it will load all the labels.
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset. Default: False.
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    num_tasks : int
+        Number of prediction tasks
+    num_labels : int
+        (DEPRECATED, use num_tasks instead) Number of prediction tasks
+
+    Raises
+    ------
+    UserWarning
+        If the raw data is changed in the remote server by the author.
+
+    Examples
+    --------
+    >>> data = QM9EdgeDataset(label_keys=['mu', 'alpha'])
+    >>> data.num_tasks
+    2
+
+    >>> # iterate over the dataset
+    >>> for graph, labels in data:
+    ...     print(graph) # get information of each graph
+    ...     print(labels) # get labels of the corresponding graph
+    ...     # your code here...
+    >>>
+    """
+
+    keys = [
+        "mu",
+        "alpha",
+        "homo",
+        "lumo",
+        "gap",
+        "r2",
+        "zpve",
+        "U0",
+        "U",
+        "H",
+        "G",
+        "Cv",
+        "U0_atom",
+        "U_atom",
+        "H_atom",
+        "G_atom",
+        "A",
+        "B",
+        "C",
+    ]
+    map_dict = {}
+
+    for i, key in enumerate(keys):
+        map_dict[key] = i
+
+    def __init__(
+        self,
+        label_keys=None,
+        raw_dir=None,
+        force_reload=False,
+        verbose=True,
+        transform=None,
+    ):
+        if label_keys is None:
+            self.label_keys = None
+            self.num_labels = 19
+        else:
+            self.label_keys = [self.map_dict[i] for i in label_keys]
+            self.num_labels = len(label_keys)
+
+        self._url = _get_dgl_url("dataset/qm9_edge.npz")
+
+        super(QM9EdgeDataset, self).__init__(
+            name="qm9Edge",
+            raw_dir=raw_dir,
+            url=self._url,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def download(self):
+        if not os.path.exists(self.npz_path):
+            download(self._url, path=self.npz_path)
+
+    def process(self):
+        self.load()
+
+    @property
+    def npz_path(self):
+        return f"{self.raw_dir}/qm9_edge.npz"
+
+    def has_cache(self):
+        return os.path.exists(self.npz_path)
+
+    def save(self):
+        np.savez_compressed(
+            self.npz_path,
+            n_node=self.n_node,
+            n_edge=self.n_edge,
+            node_attr=self.node_attr,
+            node_pos=self.node_pos,
+            edge_attr=self.edge_attr,
+            src=self.src,
+            dst=self.dst,
+            targets=self.targets,
+        )
+
+    def load(self):
+        data_dict = np.load(self.npz_path, allow_pickle=True)
+
+        self.n_node = data_dict["n_node"]
+        self.n_edge = data_dict["n_edge"]
+        self.node_attr = data_dict["node_attr"]
+        self.node_pos = data_dict["node_pos"]
+        self.edge_attr = data_dict["edge_attr"]
+        self.targets = data_dict["targets"]
+
+        self.src = data_dict["src"]
+        self.dst = data_dict["dst"]
+
+        self.n_cumsum = np.concatenate([[0], np.cumsum(self.n_node)])
+        self.ne_cumsum = np.concatenate([[0], np.cumsum(self.n_edge)])
+
+    def __getitem__(self, idx):
+        r"""Get graph and label by index
+
+        Parameters
+        ----------
+        idx : int
+            Item index
+
+        Returns
+        -------
+        dgl.DGLGraph
+           The graph contains:
+
+           - ``ndata['pos']``: the coordinates of each atom
+           - ``ndata['attr']``: the features of each atom
+           - ``edata['edge_attr']``: the features of each bond
+
+        Tensor
+            Property values of molecular graphs
+        """
+
+        pos = self.node_pos[self.n_cumsum[idx] : self.n_cumsum[idx + 1]]
+        src = self.src[self.ne_cumsum[idx] : self.ne_cumsum[idx + 1]]
+        dst = self.dst[self.ne_cumsum[idx] : self.ne_cumsum[idx + 1]]
+
+        g = dgl_graph((src, dst))
+
+        g.ndata["pos"] = F.tensor(pos, dtype=F.data_type_dict["float32"])
+        g.ndata["attr"] = F.tensor(
+            self.node_attr[self.n_cumsum[idx] : self.n_cumsum[idx + 1]],
+            dtype=F.data_type_dict["float32"],
+        )
+        g.edata["edge_attr"] = F.tensor(
+            self.edge_attr[self.ne_cumsum[idx] : self.ne_cumsum[idx + 1]],
+            dtype=F.data_type_dict["float32"],
+        )
+
+        label = F.tensor(
+            self.targets[idx][self.label_keys],
+            dtype=F.data_type_dict["float32"],
+        )
+
+        if self._transform is not None:
+            g = self._transform(g)
+
+        return g, label
+
+    def __len__(self):
+        r"""Number of graphs in the dataset.
+
+        Returns
+        -------
+        int
+        """
+        return self.n_node.shape[0]
+
+    @property
+    def num_tasks(self):
+        r"""
+        Returns
+        -------
+        int
+            Number of prediction tasks
+        """
+        return self.num_labels
+
+
+QM9Edge = QM9EdgeDataset

python/dgl/data/reddit.py

@@ -0,0 +1,223 @@
+""" Reddit dataset for community detection """
+from __future__ import absolute_import
+
+import os
+
+import numpy as np
+
+import scipy.sparse as sp
+
+from .. import backend as F
+from ..convert import from_scipy
+from ..transforms import reorder_graph
+
+from .dgl_dataset import DGLBuiltinDataset
+from .utils import (
+    _get_dgl_url,
+    deprecate_property,
+    generate_mask_tensor,
+    load_graphs,
+    save_graphs,
+)
+
+
+class RedditDataset(DGLBuiltinDataset):
+    r"""Reddit dataset for community detection (node classification)
+
+    This is a graph dataset from Reddit posts made in the month of September, 2014.
+    The node label in this case is the community, or “subreddit”, that a post belongs to.
+    The authors sampled 50 large communities and built a post-to-post graph, connecting
+    posts if the same user comments on both. In total this dataset contains 232,965
+    posts with an average degree of 492. We use the first 20 days for training and the
+    remaining days for testing (with 30% used for validation).
+
+    Reference: `<http://snap.stanford.edu/graphsage/>`_
+
+    Statistics
+
+    - Nodes: 232,965
+    - Edges: 114,615,892
+    - Node feature size: 602
+    - Number of training samples: 153,431
+    - Number of validation samples: 23,831
+    - Number of test samples: 55,703
+
+    Parameters
+    ----------
+    self_loop : bool
+        Whether load dataset with self loop connections. Default: False
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of classes for each node
+
+    Examples
+    --------
+    >>> data = RedditDataset()
+    >>> g = data[0]
+    >>> num_classes = data.num_classes
+    >>>
+    >>> # get node feature
+    >>> feat = g.ndata['feat']
+    >>>
+    >>> # get data split
+    >>> train_mask = g.ndata['train_mask']
+    >>> val_mask = g.ndata['val_mask']
+    >>> test_mask = g.ndata['test_mask']
+    >>>
+    >>> # get labels
+    >>> label = g.ndata['label']
+    >>>
+    >>> # Train, Validation and Test
+    """
+
+    def __init__(
+        self,
+        self_loop=False,
+        raw_dir=None,
+        force_reload=False,
+        verbose=False,
+        transform=None,
+    ):
+        self_loop_str = ""
+        if self_loop:
+            self_loop_str = "_self_loop"
+        _url = _get_dgl_url("dataset/reddit{}.zip".format(self_loop_str))
+        self._self_loop_str = self_loop_str
+        super(RedditDataset, self).__init__(
+            name="reddit{}".format(self_loop_str),
+            url=_url,
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        # graph
+        coo_adj = sp.load_npz(
+            os.path.join(
+                self.raw_path, "reddit{}_graph.npz".format(self._self_loop_str)
+            )
+        )
+        self._graph = from_scipy(coo_adj)
+        # features and labels
+        reddit_data = np.load(os.path.join(self.raw_path, "reddit_data.npz"))
+        features = reddit_data["feature"]
+        labels = reddit_data["label"]
+        # tarin/val/test indices
+        node_types = reddit_data["node_types"]
+        train_mask = node_types == 1
+        val_mask = node_types == 2
+        test_mask = node_types == 3
+        self._graph.ndata["train_mask"] = generate_mask_tensor(train_mask)
+        self._graph.ndata["val_mask"] = generate_mask_tensor(val_mask)
+        self._graph.ndata["test_mask"] = generate_mask_tensor(test_mask)
+        self._graph.ndata["feat"] = F.tensor(
+            features, dtype=F.data_type_dict["float32"]
+        )
+        self._graph.ndata["label"] = F.tensor(
+            labels, dtype=F.data_type_dict["int64"]
+        )
+        self._graph = reorder_graph(
+            self._graph,
+            node_permute_algo="rcmk",
+            edge_permute_algo="dst",
+            store_ids=False,
+        )
+
+        self._print_info()
+
+    def has_cache(self):
+        graph_path = os.path.join(self.save_path, "dgl_graph.bin")
+        if os.path.exists(graph_path):
+            return True
+        return False
+
+    def save(self):
+        graph_path = os.path.join(self.save_path, "dgl_graph.bin")
+        save_graphs(graph_path, self._graph)
+
+    def load(self):
+        graph_path = os.path.join(self.save_path, "dgl_graph.bin")
+        graphs, _ = load_graphs(graph_path)
+        self._graph = graphs[0]
+        self._graph.ndata["train_mask"] = generate_mask_tensor(
+            self._graph.ndata["train_mask"].numpy()
+        )
+        self._graph.ndata["val_mask"] = generate_mask_tensor(
+            self._graph.ndata["val_mask"].numpy()
+        )
+        self._graph.ndata["test_mask"] = generate_mask_tensor(
+            self._graph.ndata["test_mask"].numpy()
+        )
+        self._print_info()
+
+    def _print_info(self):
+        if self.verbose:
+            print("Finished data loading.")
+            print("  NumNodes: {}".format(self._graph.num_nodes()))
+            print("  NumEdges: {}".format(self._graph.num_edges()))
+            print("  NumFeats: {}".format(self._graph.ndata["feat"].shape[1]))
+            print("  NumClasses: {}".format(self.num_classes))
+            print(
+                "  NumTrainingSamples: {}".format(
+                    F.nonzero_1d(self._graph.ndata["train_mask"]).shape[0]
+                )
+            )
+            print(
+                "  NumValidationSamples: {}".format(
+                    F.nonzero_1d(self._graph.ndata["val_mask"]).shape[0]
+                )
+            )
+            print(
+                "  NumTestSamples: {}".format(
+                    F.nonzero_1d(self._graph.ndata["test_mask"]).shape[0]
+                )
+            )
+
+    @property
+    def num_classes(self):
+        r"""Number of classes for each node."""
+        return 41
+
+    def __getitem__(self, idx):
+        r"""Get graph by index
+
+        Parameters
+        ----------
+        idx : int
+            Item index
+
+        Returns
+        -------
+        :class:`dgl.DGLGraph`
+            graph structure, node labels, node features and splitting masks:
+
+            - ``ndata['label']``: node label
+            - ``ndata['feat']``: node feature
+            - ``ndata['train_mask']``： mask for training node set
+            - ``ndata['val_mask']``: mask for validation node set
+            - ``ndata['test_mask']:`` mask for test node set
+        """
+        assert idx == 0, "Reddit Dataset only has one graph"
+        if self._transform is None:
+            return self._graph
+        else:
+            return self._transform(self._graph)
+
+    def __len__(self):
+        r"""Number of graphs in the dataset"""
+        return 1

python/dgl/data/sbm.py

@@ -0,0 +1,276 @@
+"""Dataset for stochastic block model."""
+import math
+import os
+import random
+
+import numpy as np
+import numpy.random as npr
+import scipy as sp
+
+from .. import batch
+from ..convert import from_scipy
+from .dgl_dataset import DGLDataset
+from .utils import load_graphs, load_info, save_graphs, save_info
+
+
+def sbm(n_blocks, block_size, p, q, rng=None):
+    """(Symmetric) Stochastic Block Model
+
+    Parameters
+    ----------
+    n_blocks : int
+        Number of blocks.
+    block_size : int
+        Block size.
+    p : float
+        Probability for intra-community edge.
+    q : float
+        Probability for inter-community edge.
+    rng : numpy.random.RandomState, optional
+        Random number generator.
+
+    Returns
+    -------
+    scipy sparse matrix
+        The adjacency matrix of generated graph.
+    """
+    n = n_blocks * block_size
+    p /= n
+    q /= n
+    rng = np.random.RandomState() if rng is None else rng
+
+    rows = []
+    cols = []
+    for i in range(n_blocks):
+        for j in range(i, n_blocks):
+            density = p if i == j else q
+            block = sp.sparse.random(
+                block_size,
+                block_size,
+                density,
+                random_state=rng,
+                data_rvs=lambda n: np.ones(n),
+            )
+            rows.append(block.row + i * block_size)
+            cols.append(block.col + j * block_size)
+
+    rows = np.hstack(rows)
+    cols = np.hstack(cols)
+    a = sp.sparse.coo_matrix(
+        (np.ones(rows.shape[0]), (rows, cols)), shape=(n, n)
+    )
+    adj = sp.sparse.triu(a) + sp.sparse.triu(a, 1).transpose()
+    return adj
+
+
+class SBMMixtureDataset(DGLDataset):
+    r"""Symmetric Stochastic Block Model Mixture
+
+    Reference: Appendix C of `Supervised Community Detection with Hierarchical Graph Neural Networks <https://arxiv.org/abs/1705.08415>`_
+
+    Parameters
+    ----------
+    n_graphs : int
+        Number of graphs.
+    n_nodes : int
+        Number of nodes.
+    n_communities : int
+        Number of communities.
+    k : int, optional
+        Multiplier. Default: 2
+    avg_deg : int, optional
+        Average degree. Default: 3
+    pq : list of pair of nonnegative float or str, optional
+        Random densities. This parameter is for future extension,
+        for now it's always using the default value.
+        Default: Appendix_C
+    rng : numpy.random.RandomState, optional
+        Random number generator. If not given, it's numpy.random.RandomState() with `seed=None`,
+        which read data from /dev/urandom (or the Windows analogue) if available or seed from
+        the clock otherwise.
+        Default: None
+
+    Raises
+    ------
+    RuntimeError is raised if pq is not a list or string.
+
+    Examples
+    --------
+    >>> data = SBMMixtureDataset(n_graphs=16, n_nodes=10000, n_communities=2)
+    >>> from torch.utils.data import DataLoader
+    >>> dataloader = DataLoader(data, batch_size=1, collate_fn=data.collate_fn)
+    >>> for graph, line_graph, graph_degrees, line_graph_degrees, pm_pd in dataloader:
+    ...     # your code here
+    """
+
+    def __init__(
+        self,
+        n_graphs,
+        n_nodes,
+        n_communities,
+        k=2,
+        avg_deg=3,
+        pq="Appendix_C",
+        rng=None,
+    ):
+        self._n_graphs = n_graphs
+        self._n_nodes = n_nodes
+        self._n_communities = n_communities
+        assert n_nodes % n_communities == 0
+        self._block_size = n_nodes // n_communities
+        self._k = k
+        self._avg_deg = avg_deg
+        self._pq = pq
+        self._rng = rng
+        super(SBMMixtureDataset, self).__init__(
+            name="sbmmixture",
+            hash_key=(n_graphs, n_nodes, n_communities, k, avg_deg, pq, rng),
+        )
+
+    def process(self):
+        pq = self._pq
+        if type(pq) is list:
+            assert len(pq) == self._n_graphs
+        elif type(pq) is str:
+            generator = {"Appendix_C": self._appendix_c}[pq]
+            pq = [generator() for _ in range(self._n_graphs)]
+        else:
+            raise RuntimeError()
+        self._graphs = [
+            from_scipy(sbm(self._n_communities, self._block_size, *x))
+            for x in pq
+        ]
+        self._line_graphs = [
+            g.line_graph(backtracking=False) for g in self._graphs
+        ]
+        in_degrees = lambda g: g.in_degrees().float()
+        self._graph_degrees = [in_degrees(g) for g in self._graphs]
+        self._line_graph_degrees = [in_degrees(lg) for lg in self._line_graphs]
+        self._pm_pds = list(zip(*[g.edges() for g in self._graphs]))[0]
+
+    @property
+    def graph_path(self):
+        return os.path.join(self.save_path, "graphs_{}.bin".format(self.hash))
+
+    @property
+    def line_graph_path(self):
+        return os.path.join(
+            self.save_path, "line_graphs_{}.bin".format(self.hash)
+        )
+
+    @property
+    def info_path(self):
+        return os.path.join(self.save_path, "info_{}.pkl".format(self.hash))
+
+    def has_cache(self):
+        return (
+            os.path.exists(self.graph_path)
+            and os.path.exists(self.line_graph_path)
+            and os.path.exists(self.info_path)
+        )
+
+    def save(self):
+        save_graphs(self.graph_path, self._graphs)
+        save_graphs(self.line_graph_path, self._line_graphs)
+        save_info(
+            self.info_path,
+            {
+                "graph_degree": self._graph_degrees,
+                "line_graph_degree": self._line_graph_degrees,
+                "pm_pds": self._pm_pds,
+            },
+        )
+
+    def load(self):
+        self._graphs, _ = load_graphs(self.graph_path)
+        self._line_graphs, _ = load_graphs(self.line_graph_path)
+        info = load_info(self.info_path)
+        self._graph_degrees = info["graph_degree"]
+        self._line_graph_degrees = info["line_graph_degree"]
+        self._pm_pds = info["pm_pds"]
+
+    def __len__(self):
+        r"""Number of graphs in the dataset."""
+        return len(self._graphs)
+
+    def __getitem__(self, idx):
+        r"""Get one example by index
+
+        Parameters
+        ----------
+        idx : int
+            Item index
+
+        Returns
+        -------
+        graph: :class:`dgl.DGLGraph`
+            The original graph
+        line_graph: :class:`dgl.DGLGraph`
+            The line graph of `graph`
+        graph_degree: numpy.ndarray
+            In degrees for each node in `graph`
+        line_graph_degree: numpy.ndarray
+            In degrees for each node in `line_graph`
+        pm_pd: numpy.ndarray
+            Edge indicator matrices Pm and Pd
+        """
+        return (
+            self._graphs[idx],
+            self._line_graphs[idx],
+            self._graph_degrees[idx],
+            self._line_graph_degrees[idx],
+            self._pm_pds[idx],
+        )
+
+    def _appendix_c(self):
+        q = npr.uniform(0, self._avg_deg - math.sqrt(self._avg_deg))
+        p = self._k * self._avg_deg - q
+        if random.random() < 0.5:
+            return p, q
+        else:
+            return q, p
+
+    def collate_fn(self, x):
+        r"""The `collate` function for dataloader
+
+        Parameters
+        ----------
+        x : tuple
+            a batch of data that contains:
+
+            - graph: :class:`dgl.DGLGraph`
+                The original graph
+            - line_graph: :class:`dgl.DGLGraph`
+                The line graph of `graph`
+            - graph_degree: numpy.ndarray
+                In degrees for each node in `graph`
+            - line_graph_degree: numpy.ndarray
+                In degrees for each node in `line_graph`
+            - pm_pd: numpy.ndarray
+                Edge indicator matrices Pm and Pd
+
+        Returns
+        -------
+        g_batch: :class:`dgl.DGLGraph`
+            Batched graphs
+        lg_batch: :class:`dgl.DGLGraph`
+            Batched line graphs
+        degg_batch: numpy.ndarray
+            A batch of in degrees for each node in `g_batch`
+        deglg_batch: numpy.ndarray
+            A batch of in degrees for each node in `lg_batch`
+        pm_pd_batch: numpy.ndarray
+            A batch of edge indicator matrices Pm and Pd
+        """
+        g, lg, deg_g, deg_lg, pm_pd = zip(*x)
+        g_batch = batch.batch(g)
+        lg_batch = batch.batch(lg)
+        degg_batch = np.concatenate(deg_g, axis=0)
+        deglg_batch = np.concatenate(deg_lg, axis=0)
+        pm_pd_batch = np.concatenate(
+            [x + i * self._n_nodes for i, x in enumerate(pm_pd)], axis=0
+        )
+        return g_batch, lg_batch, degg_batch, deglg_batch, pm_pd_batch
+
+
+SBMMixture = SBMMixtureDataset

python/dgl/data/superpixel.py

@@ -0,0 +1,435 @@
+import os
+import pickle
+
+import numpy as np
+from scipy.spatial.distance import cdist
+from tqdm.auto import tqdm
+
+from .. import backend as F
+from ..convert import graph as dgl_graph
+
+from .dgl_dataset import DGLDataset
+from .utils import download, extract_archive, load_graphs, save_graphs, Subset
+
+
+def sigma(dists, kth=8):
+    num_nodes = dists.shape[0]
+
+    # Compute sigma and reshape.
+    if kth > num_nodes:
+        # Handling for graphs with num_nodes less than kth.
+        sigma = np.array([1] * num_nodes).reshape(num_nodes, 1)
+    else:
+        # Get k-nearest neighbors for each node.
+        knns = np.partition(dists, kth, axis=-1)[:, : kth + 1]
+        sigma = knns.sum(axis=1).reshape((knns.shape[0], 1)) / kth
+
+    return sigma + 1e-8
+
+
+def compute_adjacency_matrix_images(coord, feat, use_feat=True):
+    coord = coord.reshape(-1, 2)
+    # Compute coordinate distance.
+    c_dist = cdist(coord, coord)
+
+    if use_feat:
+        # Compute feature distance.
+        f_dist = cdist(feat, feat)
+        # Compute adjacency.
+        A = np.exp(
+            -((c_dist / sigma(c_dist)) ** 2) - (f_dist / sigma(f_dist)) ** 2
+        )
+    else:
+        A = np.exp(-((c_dist / sigma(c_dist)) ** 2))
+
+    # Convert to symmetric matrix.
+    A = 0.5 * (A + A.T)
+    A[np.diag_indices_from(A)] = 0
+    return A
+
+
+def compute_edges_list(A, kth=9):
+    # Get k-similar neighbor indices for each node.
+    num_nodes = A.shape[0]
+    new_kth = num_nodes - kth
+
+    if num_nodes > kth:
+        knns = np.argpartition(A, new_kth - 1, axis=-1)[:, new_kth:-1]
+        knn_values = np.partition(A, new_kth - 1, axis=-1)[:, new_kth:-1]
+    else:
+        # Handling for graphs with less than kth nodes.
+        # In such cases, the resulting graph will be fully connected.
+        knns = np.tile(np.arange(num_nodes), num_nodes).reshape(
+            num_nodes, num_nodes
+        )
+        knn_values = A
+
+        # Removing self loop.
+        if num_nodes != 1:
+            knn_values = A[knns != np.arange(num_nodes)[:, None]].reshape(
+                num_nodes, -1
+            )
+            knns = knns[knns != np.arange(num_nodes)[:, None]].reshape(
+                num_nodes, -1
+            )
+    return knns, knn_values
+
+
+class SuperPixelDataset(DGLDataset):
+    def __init__(
+        self,
+        raw_dir=None,
+        name="MNIST",
+        split="train",
+        use_feature=False,
+        force_reload=False,
+        verbose=False,
+        transform=None,
+    ):
+        assert split in ["train", "test"], "split not valid."
+        assert name in ["MNIST", "CIFAR10"], "name not valid."
+
+        self.use_feature = use_feature
+        self.split = split
+        self._dataset_name = name
+        self.graphs = []
+        self.labels = []
+
+        super().__init__(
+            name="Superpixel",
+            raw_dir=raw_dir,
+            url="""
+            https://www.dropbox.com/s/y2qwa77a0fxem47/superpixels.zip?dl=1
+            """,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    @property
+    def img_size(self):
+        r"""Size of dataset image."""
+        if self._dataset_name == "MNIST":
+            return 28
+        return 32
+
+    @property
+    def save_path(self):
+        r"""Directory to save the processed dataset."""
+        return os.path.join(self.raw_path, "processed")
+
+    @property
+    def raw_data_path(self):
+        r"""Path to save the raw dataset file."""
+        return os.path.join(self.raw_path, "superpixels.zip")
+
+    @property
+    def graph_path(self):
+        r"""Path to save the processed dataset file."""
+        if self.use_feature:
+            return os.path.join(
+                self.save_path,
+                f"use_feat_{self._dataset_name}_{self.split}.pkl",
+            )
+        return os.path.join(
+            self.save_path, f"{self._dataset_name}_{self.split}.pkl"
+        )
+
+    def download(self):
+        path = download(self.url, path=self.raw_data_path)
+        extract_archive(path, target_dir=self.raw_path, overwrite=True)
+
+    def process(self):
+        if self._dataset_name == "MNIST":
+            plk_file = "mnist_75sp"
+        elif self._dataset_name == "CIFAR10":
+            plk_file = "cifar10_150sp"
+
+        with open(
+            os.path.join(
+                self.raw_path, "superpixels", f"{plk_file}_{self.split}.pkl"
+            ),
+            "rb",
+        ) as f:
+            self.labels, self.sp_data = pickle.load(f)
+            self.labels = F.tensor(self.labels)
+
+        self.Adj_matrices = []
+        self.node_features = []
+        self.edges_lists = []
+        self.edge_features = []
+
+        for index, sample in enumerate(
+            tqdm(self.sp_data, desc=f"Processing {self.split} dataset")
+        ):
+            mean_px, coord = sample[:2]
+            coord = coord / self.img_size
+
+            if self.use_feature:
+                A = compute_adjacency_matrix_images(
+                    coord, mean_px
+                )  # using super-pixel locations + features
+            else:
+                A = compute_adjacency_matrix_images(
+                    coord, mean_px, False
+                )  # using only super-pixel locations
+            edges_list, edge_values_list = compute_edges_list(A)
+
+            N_nodes = A.shape[0]
+
+            mean_px = mean_px.reshape(N_nodes, -1)
+            coord = coord.reshape(N_nodes, 2)
+            x = np.concatenate((mean_px, coord), axis=1)
+
+            edge_values_list = edge_values_list.reshape(-1)
+
+            self.node_features.append(x)
+            self.edge_features.append(edge_values_list)
+            self.Adj_matrices.append(A)
+            self.edges_lists.append(edges_list)
+
+        for index in tqdm(
+            range(len(self.sp_data)), desc=f"Dump {self.split} dataset"
+        ):
+            N = self.node_features[index].shape[0]
+
+            src_nodes = []
+            dst_nodes = []
+            for src, dsts in enumerate(self.edges_lists[index]):
+                # handling for 1 node where the self loop would be the only edge
+                if N == 1:
+                    src_nodes.append(src)
+                    dst_nodes.append(dsts)
+                else:
+                    dsts = dsts[dsts != src]
+                    srcs = [src] * len(dsts)
+                    src_nodes.extend(srcs)
+                    dst_nodes.extend(dsts)
+
+            src_nodes = F.tensor(src_nodes)
+            dst_nodes = F.tensor(dst_nodes)
+
+            g = dgl_graph((src_nodes, dst_nodes), num_nodes=N)
+            g.ndata["feat"] = F.zerocopy_from_numpy(
+                self.node_features[index]
+            ).to(F.float32)
+            g.edata["feat"] = (
+                F.zerocopy_from_numpy(self.edge_features[index])
+                .to(F.float32)
+                .unsqueeze(1)
+            )
+
+            self.graphs.append(g)
+
+    def load(self):
+        self.graphs, label_dict = load_graphs(self.graph_path)
+        self.labels = label_dict["labels"]
+
+    def save(self):
+        save_graphs(
+            self.graph_path, self.graphs, labels={"labels": self.labels}
+        )
+
+    def has_cache(self):
+        return os.path.exists(self.graph_path)
+
+    def __len__(self):
+        return len(self.graphs)
+
+    def __getitem__(self, idx):
+        """Get the idx-th sample.
+
+        Parameters
+        ---------
+        idx : int or tensor
+            The sample index.
+            1-D tensor as `idx` is allowed when transform is None.
+
+        Returns
+        -------
+        (:class:`dgl.DGLGraph`, Tensor)
+            Graph with node feature stored in ``feat`` field and its label.
+        or
+        :class:`dgl.data.utils.Subset`
+            Subset of the dataset at specified indices
+        """
+        if F.is_tensor(idx) and idx.dim() == 1:
+            if self._transform is None:
+                return Subset(self, idx.cpu())
+
+            raise ValueError(
+                "Tensor idx not supported when transform is not None."
+            )
+
+        if self._transform is None:
+            return self.graphs[idx], self.labels[idx]
+
+        return self._transform(self.graphs[idx]), self.labels[idx]
+
+
+class MNISTSuperPixelDataset(SuperPixelDataset):
+    r"""MNIST superpixel dataset for the graph classification task.
+
+    DGL dataset of MNIST and CIFAR10 in the benchmark-gnn which contains graphs
+    converted fromt the original MINST and CIFAR10 images.
+
+    Reference `<http://arxiv.org/abs/2003.00982>`_
+
+    Statistics:
+
+        - Train examples: 60,000
+        - Test examples: 10,000
+        - Size of dataset images: 28
+
+    Parameters
+    ----------
+    raw_dir : str
+        Directory to store all the downloaded raw datasets.
+        Default: "~/.dgl/".
+    split : str
+        Should be chosen from ["train", "test"]
+        Default: "train".
+    use_feature: bool
+
+        - True: Adj matrix defined from super-pixel locations + features
+        - False: Adj matrix defined from super-pixel locations (only)
+
+        Default: False.
+    force_reload : bool
+        Whether to reload the dataset.
+        Default: False.
+    verbose : bool
+        Whether to print out progress information.
+        Default: False.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Examples
+    ---------
+    >>> from dgl.data import MNISTSuperPixelDataset
+
+    >>> # MNIST dataset
+    >>> train_dataset = MNISTSuperPixelDataset(split="train")
+    >>> len(train_dataset)
+    60000
+    >>> graph, label = train_dataset[0]
+    >>> graph
+    Graph(num_nodes=71, num_edges=568,
+        ndata_schemes={'feat': Scheme(shape=(3,), dtype=torch.float32)}
+        edata_schemes={'feat': Scheme(shape=(1,), dtype=torch.float32)})
+
+    >>> # support tensor to be index when transform is None
+    >>> # see details in __getitem__ function
+    >>> import torch
+    >>> idx = torch.tensor([0, 1, 2])
+    >>> train_dataset_subset = train_dataset[idx]
+    >>> train_dataset_subset[0]
+    Graph(num_nodes=71, num_edges=568,
+        ndata_schemes={'feat': Scheme(shape=(3,), dtype=torch.float32)}
+        edata_schemes={'feat': Scheme(shape=(1,), dtype=torch.float32)})
+    """
+
+    def __init__(
+        self,
+        raw_dir=None,
+        split="train",
+        use_feature=False,
+        force_reload=False,
+        verbose=False,
+        transform=None,
+    ):
+        super().__init__(
+            raw_dir=raw_dir,
+            name="MNIST",
+            split=split,
+            use_feature=use_feature,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+
+class CIFAR10SuperPixelDataset(SuperPixelDataset):
+    r"""CIFAR10 superpixel dataset for the graph classification task.
+
+    DGL dataset of CIFAR10 in the benchmark-gnn which contains graphs
+    converted fromt the original CIFAR10 images.
+
+    Reference `<http://arxiv.org/abs/2003.00982>`_
+
+    Statistics:
+
+        - Train examples: 50,000
+        - Test examples: 10,000
+        - Size of dataset images: 32
+
+    Parameters
+    ----------
+    raw_dir : str
+        Directory to store all the downloaded raw datasets.
+        Default: "~/.dgl/".
+    split : str
+        Should be chosen from ["train", "test"]
+        Default: "train".
+    use_feature: bool
+
+        - True: Adj matrix defined from super-pixel locations + features
+        - False: Adj matrix defined from super-pixel locations (only)
+
+        Default: False.
+    force_reload : bool
+        Whether to reload the dataset.
+        Default: False.
+    verbose : bool
+        Whether to print out progress information.
+        Default: False.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Examples
+    ---------
+    >>> from dgl.data import CIFAR10SuperPixelDataset
+
+    >>> # CIFAR10 dataset
+    >>> train_dataset = CIFAR10SuperPixelDataset(split="train")
+    >>> len(train_dataset)
+    50000
+    >>> graph, label = train_dataset[0]
+    >>> graph
+    Graph(num_nodes=123, num_edges=984,
+        ndata_schemes={'feat': Scheme(shape=(5,), dtype=torch.float32)}
+        edata_schemes={'feat': Scheme(shape=(1,), dtype=torch.float32)}),
+
+    >>> # support tensor to be index when transform is None
+    >>> # see details in __getitem__ function
+    >>> import torch
+    >>> idx = torch.tensor([0, 1, 2])
+    >>> train_dataset_subset = train_dataset[idx]
+    >>> train_dataset_subset[0]
+    Graph(num_nodes=123, num_edges=984,
+        ndata_schemes={'feat': Scheme(shape=(5,), dtype=torch.float32)}
+        edata_schemes={'feat': Scheme(shape=(1,), dtype=torch.float32)}),
+    """
+
+    def __init__(
+        self,
+        raw_dir=None,
+        split="train",
+        use_feature=False,
+        force_reload=False,
+        verbose=False,
+        transform=None,
+    ):
+        super().__init__(
+            raw_dir=raw_dir,
+            name="CIFAR10",
+            split=split,
+            use_feature=use_feature,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )

python/dgl/data/synthetic.py

@@ -0,0 +1,834 @@
+"""Synthetic graph datasets."""
+import math
+import os
+import pickle
+import random
+
+import networkx as nx
+import numpy as np
+
+from .. import backend as F
+from ..batch import batch
+from ..convert import graph
+from ..transforms import reorder_graph
+from .dgl_dataset import DGLBuiltinDataset
+from .utils import _get_dgl_url, download, load_graphs, save_graphs
+
+
+class BAShapeDataset(DGLBuiltinDataset):
+    r"""BA-SHAPES dataset from `GNNExplainer: Generating Explanations for Graph Neural Networks
+    <https://arxiv.org/abs/1903.03894>`__
+
+    This is a synthetic dataset for node classification. It is generated by performing the
+    following steps in order.
+
+    - Construct a base Barabási–Albert (BA) graph.
+    - Construct a set of five-node house-structured network motifs.
+    - Attach the motifs to randomly selected nodes of the base graph.
+    - Perturb the graph by adding random edges.
+    - Nodes are assigned to 4 classes. Nodes of label 0 belong to the base BA graph. Nodes of
+      label 1, 2, 3 are separately at the middle, bottom, or top of houses.
+    - Generate constant feature for all nodes, which is 1.
+
+    Parameters
+    ----------
+    num_base_nodes : int, optional
+        Number of nodes in the base BA graph. Default: 300
+    num_base_edges_per_node : int, optional
+        Number of edges to attach from a new node to existing nodes in constructing the base BA
+        graph. Default: 5
+    num_motifs : int, optional
+        Number of house-structured network motifs to use. Default: 80
+    perturb_ratio : float, optional
+        Number of random edges to add in perturbation divided by the number of edges in the
+        original graph. Default: 0.01
+    seed : integer, random_state, or None, optional
+        Indicator of random number generation state. Default: None
+    raw_dir : str, optional
+        Raw file directory to store the processed data. Default: ~/.dgl/
+    force_reload : bool, optional
+        Whether to always generate the data from scratch rather than load a cached version.
+        Default: False
+    verbose : bool, optional
+        Whether to print progress information. Default: True
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access. Default: None
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of node classes
+
+    Examples
+    --------
+
+    >>> from dgl.data import BAShapeDataset
+    >>> dataset = BAShapeDataset()
+    >>> dataset.num_classes
+    4
+    >>> g = dataset[0]
+    >>> label = g.ndata['label']
+    >>> feat = g.ndata['feat']
+    """
+
+    def __init__(
+        self,
+        num_base_nodes=300,
+        num_base_edges_per_node=5,
+        num_motifs=80,
+        perturb_ratio=0.01,
+        seed=None,
+        raw_dir=None,
+        force_reload=False,
+        verbose=True,
+        transform=None,
+    ):
+        self.num_base_nodes = num_base_nodes
+        self.num_base_edges_per_node = num_base_edges_per_node
+        self.num_motifs = num_motifs
+        self.perturb_ratio = perturb_ratio
+        self.seed = seed
+        super(BAShapeDataset, self).__init__(
+            name="BA-SHAPES",
+            url=None,
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        g = nx.barabasi_albert_graph(
+            self.num_base_nodes, self.num_base_edges_per_node, self.seed
+        )
+        edges = list(g.edges())
+        src, dst = map(list, zip(*edges))
+        n = self.num_base_nodes
+
+        # Nodes in the base BA graph belong to class 0
+        node_labels = [0] * n
+        # The motifs will be evenly attached to the nodes in the base graph.
+        spacing = math.floor(n / self.num_motifs)
+
+        for motif_id in range(self.num_motifs):
+            # Construct a five-node house-structured network motif
+            motif_edges = [
+                (n, n + 1),
+                (n + 1, n + 2),
+                (n + 2, n + 3),
+                (n + 3, n),
+                (n + 4, n),
+                (n + 4, n + 1),
+            ]
+            motif_src, motif_dst = map(list, zip(*motif_edges))
+            src.extend(motif_src)
+            dst.extend(motif_dst)
+
+            # Nodes at the middle of a house belong to class 1
+            # Nodes at the bottom of a house belong to class 2
+            # Nodes at the top of a house belong to class 3
+            node_labels.extend([1, 1, 2, 2, 3])
+
+            # Attach the motif to the base BA graph
+            src.append(n)
+            dst.append(int(motif_id * spacing))
+            n += 5
+
+        g = graph((src, dst), num_nodes=n)
+
+        # Perturb the graph by adding non-self-loop random edges
+        num_real_edges = g.num_edges()
+        max_ratio = (n * (n - 1) - num_real_edges) / num_real_edges
+        assert (
+            self.perturb_ratio <= max_ratio
+        ), "perturb_ratio cannot exceed {:.4f}".format(max_ratio)
+        num_random_edges = int(num_real_edges * self.perturb_ratio)
+
+        if self.seed is not None:
+            np.random.seed(self.seed)
+        for _ in range(num_random_edges):
+            while True:
+                u = np.random.randint(0, n)
+                v = np.random.randint(0, n)
+                if (not g.has_edges_between(u, v)) and (u != v):
+                    break
+            g.add_edges(u, v)
+
+        g.ndata["label"] = F.tensor(node_labels, F.int64)
+        g.ndata["feat"] = F.ones((n, 1), F.float32, F.cpu())
+        self._graph = reorder_graph(
+            g,
+            node_permute_algo="rcmk",
+            edge_permute_algo="dst",
+            store_ids=False,
+        )
+
+    @property
+    def graph_path(self):
+        return os.path.join(
+            self.save_path, "{}_dgl_graph.bin".format(self.name)
+        )
+
+    def save(self):
+        save_graphs(str(self.graph_path), self._graph)
+
+    def has_cache(self):
+        return os.path.exists(self.graph_path)
+
+    def load(self):
+        graphs, _ = load_graphs(str(self.graph_path))
+        self._graph = graphs[0]
+
+    def __getitem__(self, idx):
+        assert idx == 0, "This dataset has only one graph."
+        if self._transform is None:
+            return self._graph
+        else:
+            return self._transform(self._graph)
+
+    def __len__(self):
+        return 1
+
+    @property
+    def num_classes(self):
+        return 4
+
+
+class BACommunityDataset(DGLBuiltinDataset):
+    r"""BA-COMMUNITY dataset from `GNNExplainer: Generating Explanations for Graph Neural Networks
+    <https://arxiv.org/abs/1903.03894>`__
+
+    This is a synthetic dataset for node classification. It is generated by performing the
+    following steps in order.
+
+    - Construct a base Barabási–Albert (BA) graph.
+    - Construct a set of five-node house-structured network motifs.
+    - Attach the motifs to randomly selected nodes of the base graph.
+    - Perturb the graph by adding random edges.
+    - Nodes are assigned to 4 classes. Nodes of label 0 belong to the base BA graph. Nodes of
+      label 1, 2, 3 are separately at the middle, bottom, or top of houses.
+    - Generate normally distributed features of length 10
+    - Repeat the above steps to generate another graph. Its nodes are assigned to class
+      4, 5, 6, 7. Its node features are generated with a distinct normal distribution.
+    - Join the two graphs by randomly adding edges between them.
+
+    Parameters
+    ----------
+    num_base_nodes : int, optional
+        Number of nodes in each base BA graph. Default: 300
+    num_base_edges_per_node : int, optional
+        Number of edges to attach from a new node to existing nodes in constructing a base BA
+        graph. Default: 4
+    num_motifs : int, optional
+        Number of house-structured network motifs to use in constructing each graph. Default: 80
+    perturb_ratio : float, optional
+        Number of random edges to add to a graph in perturbation divided by the number of original
+        edges in it. Default: 0.01
+    num_inter_edges : int, optional
+        Number of random edges to add between the two graphs. Default: 350
+    seed : integer, random_state, or None, optional
+        Indicator of random number generation state. Default: None
+    raw_dir : str, optional
+        Raw file directory to store the processed data. Default: ~/.dgl/
+    force_reload : bool, optional
+        Whether to always generate the data from scratch rather than load a cached version.
+        Default: False
+    verbose : bool, optional
+        Whether to print progress information. Default: True
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access. Default: None
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of node classes
+
+    Examples
+    --------
+
+    >>> from dgl.data import BACommunityDataset
+    >>> dataset = BACommunityDataset()
+    >>> dataset.num_classes
+    8
+    >>> g = dataset[0]
+    >>> label = g.ndata['label']
+    >>> feat = g.ndata['feat']
+    """
+
+    def __init__(
+        self,
+        num_base_nodes=300,
+        num_base_edges_per_node=4,
+        num_motifs=80,
+        perturb_ratio=0.01,
+        num_inter_edges=350,
+        seed=None,
+        raw_dir=None,
+        force_reload=False,
+        verbose=True,
+        transform=None,
+    ):
+        self.num_base_nodes = num_base_nodes
+        self.num_base_edges_per_node = num_base_edges_per_node
+        self.num_motifs = num_motifs
+        self.perturb_ratio = perturb_ratio
+        self.num_inter_edges = num_inter_edges
+        self.seed = seed
+        super(BACommunityDataset, self).__init__(
+            name="BA-COMMUNITY",
+            url=None,
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        if self.seed is not None:
+            random.seed(self.seed)
+            np.random.seed(self.seed)
+
+        # Construct two BA-SHAPES graphs
+        g1 = BAShapeDataset(
+            self.num_base_nodes,
+            self.num_base_edges_per_node,
+            self.num_motifs,
+            self.perturb_ratio,
+            force_reload=True,
+            verbose=False,
+        )[0]
+        g2 = BAShapeDataset(
+            self.num_base_nodes,
+            self.num_base_edges_per_node,
+            self.num_motifs,
+            self.perturb_ratio,
+            force_reload=True,
+            verbose=False,
+        )[0]
+
+        # Join them and randomly add edges between them
+        g = batch([g1, g2])
+        num_nodes = g.num_nodes() // 2
+        src = np.random.randint(0, num_nodes, (self.num_inter_edges,))
+        dst = np.random.randint(
+            num_nodes, 2 * num_nodes, (self.num_inter_edges,)
+        )
+        src = F.astype(F.zerocopy_from_numpy(src), g.idtype)
+        dst = F.astype(F.zerocopy_from_numpy(dst), g.idtype)
+        g.add_edges(src, dst)
+        g.ndata["label"] = F.cat(
+            [g1.ndata["label"], g2.ndata["label"] + 4], dim=0
+        )
+
+        # feature generation
+        random_mu = [0.0] * 8
+        random_sigma = [1.0] * 8
+
+        mu_1, sigma_1 = np.array([-1.0] * 2 + random_mu), np.array(
+            [0.5] * 2 + random_sigma
+        )
+        feat1 = np.random.multivariate_normal(mu_1, np.diag(sigma_1), num_nodes)
+
+        mu_2, sigma_2 = np.array([1.0] * 2 + random_mu), np.array(
+            [0.5] * 2 + random_sigma
+        )
+        feat2 = np.random.multivariate_normal(mu_2, np.diag(sigma_2), num_nodes)
+
+        feat = np.concatenate([feat1, feat2])
+        g.ndata["feat"] = F.zerocopy_from_numpy(feat)
+        self._graph = reorder_graph(
+            g,
+            node_permute_algo="rcmk",
+            edge_permute_algo="dst",
+            store_ids=False,
+        )
+
+    @property
+    def graph_path(self):
+        return os.path.join(
+            self.save_path, "{}_dgl_graph.bin".format(self.name)
+        )
+
+    def save(self):
+        save_graphs(str(self.graph_path), self._graph)
+
+    def has_cache(self):
+        return os.path.exists(self.graph_path)
+
+    def load(self):
+        graphs, _ = load_graphs(str(self.graph_path))
+        self._graph = graphs[0]
+
+    def __getitem__(self, idx):
+        assert idx == 0, "This dataset has only one graph."
+        if self._transform is None:
+            return self._graph
+        else:
+            return self._transform(self._graph)
+
+    def __len__(self):
+        return 1
+
+    @property
+    def num_classes(self):
+        return 8
+
+
+class TreeCycleDataset(DGLBuiltinDataset):
+    r"""TREE-CYCLES dataset from `GNNExplainer: Generating Explanations for Graph Neural Networks
+    <https://arxiv.org/abs/1903.03894>`__
+
+    This is a synthetic dataset for node classification. It is generated by performing the
+    following steps in order.
+
+    - Construct a balanced binary tree as the base graph.
+    - Construct a set of cycle motifs.
+    - Attach the motifs to randomly selected nodes of the base graph.
+    - Perturb the graph by adding random edges.
+    - Generate constant feature for all nodes, which is 1.
+    - Nodes in the tree belong to class 0 and nodes in cycles belong to class 1.
+
+    Parameters
+    ----------
+    tree_height : int, optional
+        Height of the balanced binary tree. Default: 8
+    num_motifs : int, optional
+        Number of cycle motifs to use. Default: 60
+    cycle_size : int, optional
+        Number of nodes in a cycle motif. Default: 6
+    perturb_ratio : float, optional
+        Number of random edges to add in perturbation divided by the
+        number of original edges in the graph. Default: 0.01
+    seed : integer, random_state, or None, optional
+        Indicator of random number generation state. Default: None
+    raw_dir : str, optional
+        Raw file directory to store the processed data. Default: ~/.dgl/
+    force_reload : bool, optional
+        Whether to always generate the data from scratch rather than load a cached version.
+        Default: False
+    verbose : bool, optional
+        Whether to print progress information. Default: True
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access. Default: None
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of node classes
+
+    Examples
+    --------
+
+    >>> from dgl.data import TreeCycleDataset
+    >>> dataset = TreeCycleDataset()
+    >>> dataset.num_classes
+    2
+    >>> g = dataset[0]
+    >>> label = g.ndata['label']
+    >>> feat = g.ndata['feat']
+    """
+
+    def __init__(
+        self,
+        tree_height=8,
+        num_motifs=60,
+        cycle_size=6,
+        perturb_ratio=0.01,
+        seed=None,
+        raw_dir=None,
+        force_reload=False,
+        verbose=True,
+        transform=None,
+    ):
+        self.tree_height = tree_height
+        self.num_motifs = num_motifs
+        self.cycle_size = cycle_size
+        self.perturb_ratio = perturb_ratio
+        self.seed = seed
+        super(TreeCycleDataset, self).__init__(
+            name="TREE-CYCLES",
+            url=None,
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        if self.seed is not None:
+            np.random.seed(self.seed)
+
+        g = nx.balanced_tree(r=2, h=self.tree_height)
+        edges = list(g.edges())
+        src, dst = map(list, zip(*edges))
+        n = nx.number_of_nodes(g)
+
+        # Nodes in the base tree graph belong to class 0
+        node_labels = [0] * n
+        # The motifs will be evenly attached to the nodes in the base graph.
+        spacing = math.floor(n / self.num_motifs)
+
+        for motif_id in range(self.num_motifs):
+            # Construct a six-node cycle
+            motif_edges = [(n + i, n + i + 1) for i in range(5)]
+            motif_edges.append((n + 5, n))
+            motif_src, motif_dst = map(list, zip(*motif_edges))
+            src.extend(motif_src)
+            dst.extend(motif_dst)
+
+            # Nodes in cycles belong to class 1
+            node_labels.extend([1] * self.cycle_size)
+
+            # Attach the motif to the base tree graph
+            anchor = int(motif_id * spacing)
+            src.append(n)
+            dst.append(anchor)
+
+            if np.random.random() > 0.5:
+                a = np.random.randint(1, 4)
+                b = np.random.randint(1, 4)
+                src.append(n + a)
+                dst.append(anchor + b)
+
+            n += self.cycle_size
+
+        g = graph((src, dst), num_nodes=n)
+
+        # Perturb the graph by adding non-self-loop random edges
+        num_real_edges = g.num_edges()
+        max_ratio = (n * (n - 1) - num_real_edges) / num_real_edges
+        assert (
+            self.perturb_ratio <= max_ratio
+        ), "perturb_ratio cannot exceed {:.4f}".format(max_ratio)
+        num_random_edges = int(num_real_edges * self.perturb_ratio)
+
+        for _ in range(num_random_edges):
+            while True:
+                u = np.random.randint(0, n)
+                v = np.random.randint(0, n)
+                if (not g.has_edges_between(u, v)) and (u != v):
+                    break
+            g.add_edges(u, v)
+
+        g.ndata["label"] = F.tensor(node_labels, F.int64)
+        g.ndata["feat"] = F.ones((n, 1), F.float32, F.cpu())
+        self._graph = reorder_graph(
+            g,
+            node_permute_algo="rcmk",
+            edge_permute_algo="dst",
+            store_ids=False,
+        )
+
+    @property
+    def graph_path(self):
+        return os.path.join(
+            self.save_path, "{}_dgl_graph.bin".format(self.name)
+        )
+
+    def save(self):
+        save_graphs(str(self.graph_path), self._graph)
+
+    def has_cache(self):
+        return os.path.exists(self.graph_path)
+
+    def load(self):
+        graphs, _ = load_graphs(str(self.graph_path))
+        self._graph = graphs[0]
+
+    def __getitem__(self, idx):
+        assert idx == 0, "This dataset has only one graph."
+        if self._transform is None:
+            return self._graph
+        else:
+            return self._transform(self._graph)
+
+    def __len__(self):
+        return 1
+
+    @property
+    def num_classes(self):
+        return 2
+
+
+class TreeGridDataset(DGLBuiltinDataset):
+    r"""TREE-GRIDS dataset from `GNNExplainer: Generating Explanations for Graph Neural Networks
+    <https://arxiv.org/abs/1903.03894>`__
+
+    This is a synthetic dataset for node classification. It is generated by performing the
+    following steps in order.
+
+    - Construct a balanced binary tree as the base graph.
+    - Construct a set of n-by-n grid motifs.
+    - Attach the motifs to randomly selected nodes of the base graph.
+    - Perturb the graph by adding random edges.
+    - Generate constant feature for all nodes, which is 1.
+    - Nodes in the tree belong to class 0 and nodes in grids belong to class 1.
+
+    Parameters
+    ----------
+    tree_height : int, optional
+        Height of the balanced binary tree. Default: 8
+    num_motifs : int, optional
+        Number of grid motifs to use. Default: 80
+    grid_size : int, optional
+        The number of nodes in a grid motif will be grid_size ^ 2. Default: 3
+    perturb_ratio : float, optional
+        Number of random edges to add in perturbation divided by the
+        number of original edges in the graph. Default: 0.1
+    seed : integer, random_state, or None, optional
+        Indicator of random number generation state. Default: None
+    raw_dir : str, optional
+        Raw file directory to store the processed data. Default: ~/.dgl/
+    force_reload : bool, optional
+        Whether to always generate the data from scratch rather than load a cached version.
+        Default: False
+    verbose : bool, optional
+        Whether to print progress information. Default: True
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access. Default: None
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of node classes
+
+    Examples
+    --------
+
+    >>> from dgl.data import TreeGridDataset
+    >>> dataset = TreeGridDataset()
+    >>> dataset.num_classes
+    2
+    >>> g = dataset[0]
+    >>> label = g.ndata['label']
+    >>> feat = g.ndata['feat']
+    """
+
+    def __init__(
+        self,
+        tree_height=8,
+        num_motifs=80,
+        grid_size=3,
+        perturb_ratio=0.1,
+        seed=None,
+        raw_dir=None,
+        force_reload=False,
+        verbose=True,
+        transform=None,
+    ):
+        self.tree_height = tree_height
+        self.num_motifs = num_motifs
+        self.grid_size = grid_size
+        self.perturb_ratio = perturb_ratio
+        self.seed = seed
+        super(TreeGridDataset, self).__init__(
+            name="TREE-GRIDS",
+            url=None,
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        if self.seed is not None:
+            np.random.seed(self.seed)
+
+        g = nx.balanced_tree(r=2, h=self.tree_height)
+        edges = list(g.edges())
+        src, dst = map(list, zip(*edges))
+        n = nx.number_of_nodes(g)
+
+        # Nodes in the base tree graph belong to class 0
+        node_labels = [0] * n
+        # The motifs will be evenly attached to the nodes in the base graph.
+        spacing = math.floor(n / self.num_motifs)
+
+        # Construct an n-by-n grid
+        motif_g = nx.grid_graph([self.grid_size, self.grid_size])
+        grid_size = nx.number_of_nodes(motif_g)
+        motif_g = nx.convert_node_labels_to_integers(motif_g, first_label=0)
+        motif_edges = list(motif_g.edges())
+        motif_src, motif_dst = map(list, zip(*motif_edges))
+        motif_src, motif_dst = np.array(motif_src), np.array(motif_dst)
+
+        for motif_id in range(self.num_motifs):
+            src.extend((motif_src + n).tolist())
+            dst.extend((motif_dst + n).tolist())
+
+            # Nodes in grids belong to class 1
+            node_labels.extend([1] * grid_size)
+
+            # Attach the motif to the base tree graph
+            src.append(n)
+            dst.append(int(motif_id * spacing))
+
+            n += grid_size
+
+        g = graph((src, dst), num_nodes=n)
+
+        # Perturb the graph by adding non-self-loop random edges
+        num_real_edges = g.num_edges()
+        max_ratio = (n * (n - 1) - num_real_edges) / num_real_edges
+        assert (
+            self.perturb_ratio <= max_ratio
+        ), "perturb_ratio cannot exceed {:.4f}".format(max_ratio)
+        num_random_edges = int(num_real_edges * self.perturb_ratio)
+
+        for _ in range(num_random_edges):
+            while True:
+                u = np.random.randint(0, n)
+                v = np.random.randint(0, n)
+                if (not g.has_edges_between(u, v)) and (u != v):
+                    break
+            g.add_edges(u, v)
+
+        g.ndata["label"] = F.tensor(node_labels, F.int64)
+        g.ndata["feat"] = F.ones((n, 1), F.float32, F.cpu())
+        self._graph = reorder_graph(
+            g,
+            node_permute_algo="rcmk",
+            edge_permute_algo="dst",
+            store_ids=False,
+        )
+
+    @property
+    def graph_path(self):
+        return os.path.join(
+            self.save_path, "{}_dgl_graph.bin".format(self.name)
+        )
+
+    def save(self):
+        save_graphs(str(self.graph_path), self._graph)
+
+    def has_cache(self):
+        return os.path.exists(self.graph_path)
+
+    def load(self):
+        graphs, _ = load_graphs(str(self.graph_path))
+        self._graph = graphs[0]
+
+    def __getitem__(self, idx):
+        assert idx == 0, "This dataset has only one graph."
+        if self._transform is None:
+            return self._graph
+        else:
+            return self._transform(self._graph)
+
+    def __len__(self):
+        return 1
+
+    @property
+    def num_classes(self):
+        return 2
+
+
+class BA2MotifDataset(DGLBuiltinDataset):
+    r"""BA-2motifs dataset from `Parameterized Explainer for Graph Neural Network
+    <https://arxiv.org/abs/2011.04573>`__
+
+    This is a synthetic dataset for graph classification. It was generated by
+    performing the following steps in order.
+
+    - Construct 1000 base Barabási–Albert (BA) graphs.
+    - Attach house-structured network motifs to half of the base BA graphs.
+    - Attach five-node cycle motifs to the rest base BA graphs.
+    - Assign each graph to one of two classes according to the type of the attached motif.
+
+    Parameters
+    ----------
+    raw_dir : str, optional
+        Raw file directory to download and store the data. Default: ~/.dgl/
+    force_reload : bool, optional
+        Whether to reload the dataset. Default: False
+    verbose : bool, optional
+        Whether to print progress information. Default: True
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access. Default: None
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of graph classes
+
+    Examples
+    --------
+
+    >>> from dgl.data import BA2MotifDataset
+    >>> dataset = BA2MotifDataset()
+    >>> dataset.num_classes
+    2
+    >>> # Get the first graph and its label
+    >>> g, label = dataset[0]
+    >>> feat = g.ndata['feat']
+    """
+
+    def __init__(
+        self, raw_dir=None, force_reload=False, verbose=True, transform=None
+    ):
+        super(BA2MotifDataset, self).__init__(
+            name="BA-2motifs",
+            url=_get_dgl_url("dataset/BA-2motif.pkl"),
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def download(self):
+        r"""Automatically download data."""
+        file_path = os.path.join(self.raw_dir, self.name + ".pkl")
+        download(self.url, path=file_path)
+
+    def process(self):
+        file_path = os.path.join(self.raw_dir, self.name + ".pkl")
+        with open(file_path, "rb") as f:
+            adjs, features, labels = pickle.load(f)
+
+        self.graphs = []
+        self.labels = F.tensor(labels, F.int64)
+
+        for i in range(len(adjs)):
+            g = graph(adjs[i].nonzero())
+            g.ndata["feat"] = F.zerocopy_from_numpy(features[i])
+            self.graphs.append(g)
+
+    @property
+    def graph_path(self):
+        return os.path.join(
+            self.save_path, "{}_dgl_graph.bin".format(self.name)
+        )
+
+    def save(self):
+        label_dict = {"labels": self.labels}
+        save_graphs(str(self.graph_path), self.graphs, label_dict)
+
+    def has_cache(self):
+        return os.path.exists(self.graph_path)
+
+    def load(self):
+        self.graphs, label_dict = load_graphs(str(self.graph_path))
+        self.labels = label_dict["labels"]
+
+    def __getitem__(self, idx):
+        g = self.graphs[idx]
+        if self._transform is not None:
+            g = self._transform(g)
+        return g, self.labels[idx]
+
+    def __len__(self):
+        return len(self.graphs)
+
+    @property
+    def num_classes(self):
+        return 2

python/dgl/data/tensor_serialize.py

@@ -0,0 +1,69 @@
+"""For Tensor Serialization"""
+from __future__ import absolute_import
+
+from .. import backend as F
+from .._ffi.function import _init_api
+from ..ndarray import NDArray
+
+__all__ = ["save_tensors", "load_tensors"]
+
+_init_api("dgl.data.tensor_serialize")
+
+
+def save_tensors(filename, tensor_dict):
+    """
+    Save dict of tensors to file
+
+    Parameters
+    ----------
+    filename : str
+        File name to store dict of tensors.
+    tensor_dict: dict of dgl NDArray or backend tensor
+        Python dict using string as key and tensor as value
+
+    Returns
+    ----------
+    status : bool
+        Return whether save operation succeeds
+    """
+    nd_dict = {}
+    is_empty_dict = len(tensor_dict) == 0
+    for key, value in tensor_dict.items():
+        if not isinstance(key, str):
+            raise Exception("Dict key has to be str")
+        if F.is_tensor(value):
+            nd_dict[key] = F.zerocopy_to_dgl_ndarray(value)
+        elif isinstance(value, NDArray):
+            nd_dict[key] = value
+        else:
+            raise Exception(
+                "Dict value has to be backend tensor or dgl ndarray"
+            )
+
+    return _CAPI_SaveNDArrayDict(filename, nd_dict, is_empty_dict)
+
+
+def load_tensors(filename, return_dgl_ndarray=False):
+    """
+    load dict of tensors from file
+
+    Parameters
+    ----------
+    filename : str
+        File name to load dict of tensors.
+    return_dgl_ndarray: bool
+        Whether return dict of dgl NDArrays or backend tensors
+
+    Returns
+    ---------
+    tensor_dict : dict
+        dict of tensor or ndarray based on return_dgl_ndarray flag
+    """
+    nd_dict = _CAPI_LoadNDArrayDict(filename)
+    tensor_dict = {}
+    for key, value in nd_dict.items():
+        if return_dgl_ndarray:
+            tensor_dict[key] = value
+        else:
+            tensor_dict[key] = F.zerocopy_from_dgl_ndarray(value)
+    return tensor_dict

python/dgl/data/tree.py

@@ -0,0 +1,305 @@
+"""Tree-structured data.
+Including:
+    - Stanford Sentiment Treebank
+"""
+from __future__ import absolute_import
+
+import os
+
+from collections import OrderedDict
+
+import networkx as nx
+
+import numpy as np
+
+from .. import backend as F
+from ..convert import from_networkx
+
+from .dgl_dataset import DGLBuiltinDataset
+from .utils import (
+    _get_dgl_url,
+    deprecate_property,
+    load_graphs,
+    load_info,
+    save_graphs,
+    save_info,
+)
+
+__all__ = ["SST", "SSTDataset"]
+
+
+class SSTDataset(DGLBuiltinDataset):
+    r"""Stanford Sentiment Treebank dataset.
+
+    Each sample is the constituency tree of a sentence. The leaf nodes
+    represent words. The word is a int value stored in the ``x`` feature field.
+    The non-leaf node has a special value ``PAD_WORD`` in the ``x`` field.
+    Each node also has a sentiment annotation: 5 classes (very negative,
+    negative, neutral, positive and very positive). The sentiment label is a
+    int value stored in the ``y`` feature field.
+    Official site: `<http://nlp.stanford.edu/sentiment/index.html>`_
+
+    Statistics:
+
+    - Train examples: 8,544
+    - Dev examples: 1,101
+    - Test examples: 2,210
+    - Number of classes for each node: 5
+
+    Parameters
+    ----------
+    mode : str, optional
+        Should be one of ['train', 'dev', 'test', 'tiny']
+        Default: train
+    glove_embed_file : str, optional
+        The path to pretrained glove embedding file.
+        Default: None
+    vocab_file : str, optional
+        Optional vocabulary file. If not given, the default vacabulary file is used.
+        Default: None
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    vocab : OrderedDict
+        Vocabulary of the dataset
+    num_classes : int
+        Number of classes for each node
+    pretrained_emb: Tensor
+        Pretrained glove embedding with respect the vocabulary.
+    vocab_size : int
+        The size of the vocabulary
+
+    Notes
+    -----
+    All the samples will be loaded and preprocessed in the memory first.
+
+    Examples
+    --------
+    >>> # get dataset
+    >>> train_data = SSTDataset()
+    >>> dev_data = SSTDataset(mode='dev')
+    >>> test_data = SSTDataset(mode='test')
+    >>> tiny_data = SSTDataset(mode='tiny')
+    >>>
+    >>> len(train_data)
+    8544
+    >>> train_data.num_classes
+    5
+    >>> glove_embed = train_data.pretrained_emb
+    >>> train_data.vocab_size
+    19536
+    >>> train_data[0]
+    Graph(num_nodes=71, num_edges=70,
+      ndata_schemes={'x': Scheme(shape=(), dtype=torch.int64), 'y': Scheme(shape=(), dtype=torch.int64), 'mask': Scheme(shape=(), dtype=torch.int64)}
+      edata_schemes={})
+    >>> for tree in train_data:
+    ...     input_ids = tree.ndata['x']
+    ...     labels = tree.ndata['y']
+    ...     mask = tree.ndata['mask']
+    ...     # your code here
+    """
+
+    PAD_WORD = -1  # special pad word id
+    UNK_WORD = -1  # out-of-vocabulary word id
+
+    def __init__(
+        self,
+        mode="train",
+        glove_embed_file=None,
+        vocab_file=None,
+        raw_dir=None,
+        force_reload=False,
+        verbose=False,
+        transform=None,
+    ):
+        assert mode in ["train", "dev", "test", "tiny"]
+        _url = _get_dgl_url("dataset/sst.zip")
+        self._glove_embed_file = glove_embed_file if mode == "train" else None
+        self.mode = mode
+        self._vocab_file = vocab_file
+        super(SSTDataset, self).__init__(
+            name="sst",
+            url=_url,
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        from nltk.corpus.reader import BracketParseCorpusReader
+
+        # load vocab file
+        self._vocab = OrderedDict()
+        vocab_file = (
+            self._vocab_file
+            if self._vocab_file is not None
+            else os.path.join(self.raw_path, "vocab.txt")
+        )
+        with open(vocab_file, encoding="utf-8") as vf:
+            for line in vf.readlines():
+                line = line.strip()
+                self._vocab[line] = len(self._vocab)
+
+        # filter glove
+        if self._glove_embed_file is not None and os.path.exists(
+            self._glove_embed_file
+        ):
+            glove_emb = {}
+            with open(self._glove_embed_file, "r", encoding="utf-8") as pf:
+                for line in pf.readlines():
+                    sp = line.split(" ")
+                    if sp[0].lower() in self._vocab:
+                        glove_emb[sp[0].lower()] = np.asarray(
+                            [float(x) for x in sp[1:]]
+                        )
+        files = ["{}.txt".format(self.mode)]
+        corpus = BracketParseCorpusReader(self.raw_path, files)
+        sents = corpus.parsed_sents(files[0])
+
+        # initialize with glove
+        pretrained_emb = []
+        fail_cnt = 0
+        for line in self._vocab.keys():
+            if self._glove_embed_file is not None and os.path.exists(
+                self._glove_embed_file
+            ):
+                if not line.lower() in glove_emb:
+                    fail_cnt += 1
+                pretrained_emb.append(
+                    glove_emb.get(
+                        line.lower(), np.random.uniform(-0.05, 0.05, 300)
+                    )
+                )
+
+        self._pretrained_emb = None
+        if self._glove_embed_file is not None and os.path.exists(
+            self._glove_embed_file
+        ):
+            self._pretrained_emb = F.tensor(np.stack(pretrained_emb, 0))
+            print(
+                "Miss word in GloVe {0:.4f}".format(
+                    1.0 * fail_cnt / len(self._pretrained_emb)
+                )
+            )
+        # build trees
+        self._trees = []
+        for sent in sents:
+            self._trees.append(self._build_tree(sent))
+
+    def _build_tree(self, root):
+        g = nx.DiGraph()
+
+        def _rec_build(nid, node):
+            for child in node:
+                cid = g.number_of_nodes()
+                if isinstance(child[0], str) or isinstance(child[0], bytes):
+                    # leaf node
+                    word = self.vocab.get(child[0].lower(), self.UNK_WORD)
+                    g.add_node(cid, x=word, y=int(child.label()), mask=1)
+                else:
+                    g.add_node(
+                        cid, x=SSTDataset.PAD_WORD, y=int(child.label()), mask=0
+                    )
+                    _rec_build(cid, child)
+                g.add_edge(cid, nid)
+
+        # add root
+        g.add_node(0, x=SSTDataset.PAD_WORD, y=int(root.label()), mask=0)
+        _rec_build(0, root)
+        ret = from_networkx(g, node_attrs=["x", "y", "mask"])
+        return ret
+
+    @property
+    def graph_path(self):
+        return os.path.join(self.save_path, self.mode + "_dgl_graph.bin")
+
+    @property
+    def vocab_path(self):
+        return os.path.join(self.save_path, "vocab.pkl")
+
+    def has_cache(self):
+        return os.path.exists(self.graph_path) and os.path.exists(
+            self.vocab_path
+        )
+
+    def save(self):
+        save_graphs(self.graph_path, self._trees)
+        save_info(self.vocab_path, {"vocab": self.vocab})
+        if self.pretrained_emb:
+            emb_path = os.path.join(self.save_path, "emb.pkl")
+            save_info(emb_path, {"embed": self.pretrained_emb})
+
+    def load(self):
+        emb_path = os.path.join(self.save_path, "emb.pkl")
+
+        self._trees = load_graphs(self.graph_path)[0]
+        self._vocab = load_info(self.vocab_path)["vocab"]
+        self._pretrained_emb = None
+        if os.path.exists(emb_path):
+            self._pretrained_emb = load_info(emb_path)["embed"]
+
+    @property
+    def vocab(self):
+        r"""Vocabulary
+
+        Returns
+        -------
+        OrderedDict
+        """
+        return self._vocab
+
+    @property
+    def pretrained_emb(self):
+        r"""Pre-trained word embedding, if given."""
+        return self._pretrained_emb
+
+    def __getitem__(self, idx):
+        r"""Get graph by index
+
+        Parameters
+        ----------
+        idx : int
+
+        Returns
+        -------
+        :class:`dgl.DGLGraph`
+
+            graph structure, word id for each node, node labels and masks.
+
+            - ``ndata['x']``: word id of the node
+            - ``ndata['y']:`` label of the node
+            - ``ndata['mask']``: 1 if the node is a leaf, otherwise 0
+        """
+        if self._transform is None:
+            return self._trees[idx]
+        else:
+            return self._transform(self._trees[idx])
+
+    def __len__(self):
+        r"""Number of graphs in the dataset."""
+        return len(self._trees)
+
+    @property
+    def vocab_size(self):
+        r"""Vocabulary size."""
+        return len(self._vocab)
+
+    @property
+    def num_classes(self):
+        r"""Number of classes for each node."""
+        return 5
+
+
+SST = SSTDataset

python/dgl/data/tu.py

@@ -0,0 +1,532 @@
+from __future__ import absolute_import
+
+import os
+
+import numpy as np
+
+from .. import backend as F
+from ..convert import graph as dgl_graph
+
+from .dgl_dataset import DGLBuiltinDataset
+from .utils import load_graphs, load_info, loadtxt, save_graphs, save_info
+
+
+class LegacyTUDataset(DGLBuiltinDataset):
+    r"""LegacyTUDataset contains lots of graph kernel datasets for graph classification.
+
+    Parameters
+    ----------
+    name : str
+        Dataset Name, such as ``ENZYMES``, ``DD``, ``COLLAB``, ``MUTAG``, can be the
+        datasets name on `<https://chrsmrrs.github.io/datasets/docs/datasets/>`_.
+    use_pandas : bool
+        Numpy's file read function has performance issue when file is large,
+        using pandas can be faster.
+        Default: False
+    hidden_size : int
+        Some dataset doesn't contain features.
+        Use constant node features initialization instead, with hidden size as ``hidden_size``.
+        Default : 10
+    max_allow_node : int
+        Remove graphs that contains more nodes than ``max_allow_node``.
+        Default : None
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    max_num_node : int
+        Maximum number of nodes
+    num_classes : int
+        Number of classes
+    num_labels : numpy.int64
+        (DEPRECATED, use num_classes instead) Number of classes
+
+    Notes
+    -----
+    LegacyTUDataset uses provided node feature by default. If no feature provided, it uses one-hot node label instead.
+    If neither labels provided, it uses constant for node feature.
+
+    The dataset sorts graphs by their labels.
+    Shuffle is preferred before manual train/val split.
+
+    Examples
+    --------
+    >>> data = LegacyTUDataset('DD')
+
+    The dataset instance is an iterable
+
+    >>> len(data)
+    1178
+    >>> g, label = data[1024]
+    >>> g
+    Graph(num_nodes=88, num_edges=410,
+          ndata_schemes={'feat': Scheme(shape=(89,), dtype=torch.float32), '_ID': Scheme(shape=(), dtype=torch.int64)}
+          edata_schemes={'_ID': Scheme(shape=(), dtype=torch.int64)})
+    >>> label
+    tensor(1)
+
+    Batch the graphs and labels for mini-batch training
+
+    >>> graphs, labels = zip(*[data[i] for i in range(16)])
+    >>> batched_graphs = dgl.batch(graphs)
+    >>> batched_labels = torch.tensor(labels)
+    >>> batched_graphs
+    Graph(num_nodes=9539, num_edges=47382,
+          ndata_schemes={'feat': Scheme(shape=(89,), dtype=torch.float32), '_ID': Scheme(shape=(), dtype=torch.int64)}
+          edata_schemes={'_ID': Scheme(shape=(), dtype=torch.int64)})
+    """
+
+    _url = r"https://www.chrsmrrs.com/graphkerneldatasets/{}.zip"
+
+    def __init__(
+        self,
+        name,
+        use_pandas=False,
+        hidden_size=10,
+        max_allow_node=None,
+        raw_dir=None,
+        force_reload=False,
+        verbose=False,
+        transform=None,
+    ):
+        url = self._url.format(name)
+        self.hidden_size = hidden_size
+        self.max_allow_node = max_allow_node
+        self.use_pandas = use_pandas
+        super(LegacyTUDataset, self).__init__(
+            name=name,
+            url=url,
+            raw_dir=raw_dir,
+            hash_key=(name, use_pandas, hidden_size, max_allow_node),
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        self.data_mode = None
+
+        if self.use_pandas:
+            import pandas as pd
+
+            DS_edge_list = self._idx_from_zero(
+                pd.read_csv(
+                    self._file_path("A"), delimiter=",", dtype=int, header=None
+                ).values
+            )
+        else:
+            DS_edge_list = self._idx_from_zero(
+                np.genfromtxt(self._file_path("A"), delimiter=",", dtype=int)
+            )
+
+        DS_indicator = self._idx_from_zero(
+            np.genfromtxt(self._file_path("graph_indicator"), dtype=int)
+        )
+        if os.path.exists(self._file_path("graph_labels")):
+            DS_graph_labels = self._idx_from_zero(
+                np.genfromtxt(self._file_path("graph_labels"), dtype=int)
+            )
+            self.num_labels = max(DS_graph_labels) + 1
+            self.graph_labels = DS_graph_labels
+        elif os.path.exists(self._file_path("graph_attributes")):
+            DS_graph_labels = np.genfromtxt(
+                self._file_path("graph_attributes"), dtype=float
+            )
+            self.num_labels = None
+            self.graph_labels = DS_graph_labels
+        else:
+            raise Exception("Unknown graph label or graph attributes")
+
+        g = dgl_graph(([], []))
+        g.add_nodes(int(DS_edge_list.max()) + 1)
+        g.add_edges(DS_edge_list[:, 0], DS_edge_list[:, 1])
+
+        node_idx_list = []
+        self.max_num_node = 0
+        for idx in range(np.max(DS_indicator) + 1):
+            node_idx = np.where(DS_indicator == idx)
+            node_idx_list.append(node_idx[0])
+            if len(node_idx[0]) > self.max_num_node:
+                self.max_num_node = len(node_idx[0])
+
+        self.graph_lists = [g.subgraph(node_idx) for node_idx in node_idx_list]
+
+        try:
+            DS_node_labels = self._idx_from_zero(
+                np.loadtxt(self._file_path("node_labels"), dtype=int)
+            )
+            g.ndata["node_label"] = F.tensor(DS_node_labels)
+            one_hot_node_labels = self._to_onehot(DS_node_labels)
+            for idxs, g in zip(node_idx_list, self.graph_lists):
+                g.ndata["feat"] = F.tensor(
+                    one_hot_node_labels[idxs, :], F.float32
+                )
+            self.data_mode = "node_label"
+        except IOError:
+            print("No Node Label Data")
+
+        try:
+            DS_node_attr = np.loadtxt(
+                self._file_path("node_attributes"), delimiter=","
+            )
+            if DS_node_attr.ndim == 1:
+                DS_node_attr = np.expand_dims(DS_node_attr, -1)
+            for idxs, g in zip(node_idx_list, self.graph_lists):
+                g.ndata["feat"] = F.tensor(DS_node_attr[idxs, :], F.float32)
+            self.data_mode = "node_attr"
+        except IOError:
+            print("No Node Attribute Data")
+
+        if "feat" not in g.ndata.keys():
+            for idxs, g in zip(node_idx_list, self.graph_lists):
+                g.ndata["feat"] = F.ones(
+                    (g.num_nodes(), self.hidden_size), F.float32, F.cpu()
+                )
+            self.data_mode = "constant"
+            if self.verbose:
+                print(
+                    "Use Constant one as Feature with hidden size {}".format(
+                        self.hidden_size
+                    )
+                )
+
+        # remove graphs that are too large by user given standard
+        # optional pre-processing steop in conformity with Rex Ying's original
+        # DiffPool implementation
+        if self.max_allow_node:
+            preserve_idx = []
+            if self.verbose:
+                print("original dataset length : ", len(self.graph_lists))
+            for i, g in enumerate(self.graph_lists):
+                if g.num_nodes() <= self.max_allow_node:
+                    preserve_idx.append(i)
+            self.graph_lists = [self.graph_lists[i] for i in preserve_idx]
+            if self.verbose:
+                print(
+                    "after pruning graphs that are too big : ",
+                    len(self.graph_lists),
+                )
+            self.graph_labels = [self.graph_labels[i] for i in preserve_idx]
+            self.max_num_node = self.max_allow_node
+        self.graph_labels = F.tensor(self.graph_labels)
+
+    def save(self):
+        label_dict = {"labels": self.graph_labels}
+        info_dict = {
+            "max_num_node": self.max_num_node,
+            "num_labels": self.num_labels,
+        }
+        save_graphs(str(self.graph_path), self.graph_lists, label_dict)
+        save_info(str(self.info_path), info_dict)
+
+    def load(self):
+        graphs, label_dict = load_graphs(str(self.graph_path))
+        info_dict = load_info(str(self.info_path))
+
+        self.graph_lists = graphs
+        self.graph_labels = label_dict["labels"]
+        self.max_num_node = info_dict["max_num_node"]
+        self.num_labels = info_dict["num_labels"]
+
+    @property
+    def graph_path(self):
+        return os.path.join(
+            self.save_path, "legacy_tu_{}_{}.bin".format(self.name, self.hash)
+        )
+
+    @property
+    def info_path(self):
+        return os.path.join(
+            self.save_path, "legacy_tu_{}_{}.pkl".format(self.name, self.hash)
+        )
+
+    def has_cache(self):
+        if os.path.exists(self.graph_path) and os.path.exists(self.info_path):
+            return True
+        return False
+
+    def __getitem__(self, idx):
+        """Get the idx-th sample.
+
+        Parameters
+        ---------
+        idx : int
+            The sample index.
+
+        Returns
+        -------
+        (:class:`dgl.DGLGraph`, Tensor)
+            Graph with node feature stored in ``feat`` field and node label in ``node_label`` if available.
+            And its label.
+        """
+        g = self.graph_lists[idx]
+        if self._transform is not None:
+            g = self._transform(g)
+        return g, self.graph_labels[idx]
+
+    def __len__(self):
+        """Return the number of graphs in the dataset."""
+        return len(self.graph_lists)
+
+    def _file_path(self, category):
+        return os.path.join(
+            self.raw_path, self.name, "{}_{}.txt".format(self.name, category)
+        )
+
+    @staticmethod
+    def _idx_from_zero(idx_tensor):
+        return idx_tensor - np.min(idx_tensor)
+
+    @staticmethod
+    def _to_onehot(label_tensor):
+        label_num = label_tensor.shape[0]
+        assert np.min(label_tensor) == 0
+        one_hot_tensor = np.zeros((label_num, np.max(label_tensor) + 1))
+        one_hot_tensor[np.arange(label_num), label_tensor] = 1
+        return one_hot_tensor
+
+    def statistics(self):
+        return (
+            self.graph_lists[0].ndata["feat"].shape[1],
+            self.num_labels,
+            self.max_num_node,
+        )
+
+    @property
+    def num_classes(self):
+        return int(self.num_labels)
+
+
+class TUDataset(DGLBuiltinDataset):
+    r"""
+    TUDataset contains lots of graph kernel datasets for graph classification.
+
+    Parameters
+    ----------
+    name : str
+        Dataset Name, such as ``ENZYMES``, ``DD``, ``COLLAB``, ``MUTAG``, can be the
+        datasets name on `<https://chrsmrrs.github.io/datasets/docs/datasets/>`_.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    max_num_node : int
+        Maximum number of nodes
+    num_classes : int
+        Number of classes
+    num_labels : int
+        (DEPRECATED, use num_classes instead) Number of classes
+
+    Notes
+    -----
+    **IMPORTANT:** Some of the datasets have duplicate edges exist in the graphs, e.g.
+    the edges in ``IMDB-BINARY`` are all duplicated.  DGL faithfully keeps the duplicates
+    as per the original data.  Other frameworks such as PyTorch Geometric removes the
+    duplicates by default.  You can remove the duplicate edges with :func:`dgl.to_simple`.
+
+    Graphs may have node labels, node attributes, edge labels, and edge attributes,
+    varing from different dataset.
+
+    Labels are mapped to :math:`\lbrace 0,\cdots,n-1 \rbrace` where :math:`n` is the
+    number of labels (some datasets have raw labels :math:`\lbrace -1, 1 \rbrace` which
+    will be mapped to :math:`\lbrace 0, 1 \rbrace`). In previous versions, the minimum
+    label was added so that :math:`\lbrace -1, 1 \rbrace` was mapped to
+    :math:`\lbrace 0, 2 \rbrace`.
+
+    The dataset sorts graphs by their labels.
+    Shuffle is preferred before manual train/val split.
+
+    Examples
+    --------
+    >>> data = TUDataset('DD')
+
+    The dataset instance is an iterable
+
+    >>> len(data)
+    1178
+    >>> g, label = data[1024]
+    >>> g
+    Graph(num_nodes=88, num_edges=410,
+          ndata_schemes={'_ID': Scheme(shape=(), dtype=torch.int64), 'node_labels': Scheme(shape=(1,), dtype=torch.int64)}
+          edata_schemes={'_ID': Scheme(shape=(), dtype=torch.int64)})
+    >>> label
+    tensor([1])
+
+    Batch the graphs and labels for mini-batch training
+
+    >>> graphs, labels = zip(*[data[i] for i in range(16)])
+    >>> batched_graphs = dgl.batch(graphs)
+    >>> batched_labels = torch.tensor(labels)
+    >>> batched_graphs
+    Graph(num_nodes=9539, num_edges=47382,
+          ndata_schemes={'node_labels': Scheme(shape=(1,), dtype=torch.int64), '_ID': Scheme(shape=(), dtype=torch.int64)}
+          edata_schemes={'_ID': Scheme(shape=(), dtype=torch.int64)})
+
+    """
+
+    _url = r"https://www.chrsmrrs.com/graphkerneldatasets/{}.zip"
+
+    def __init__(
+        self,
+        name,
+        raw_dir=None,
+        force_reload=False,
+        verbose=False,
+        transform=None,
+    ):
+        url = self._url.format(name)
+        super(TUDataset, self).__init__(
+            name=name,
+            url=url,
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        DS_edge_list = self._idx_from_zero(
+            loadtxt(self._file_path("A"), delimiter=",").astype(int)
+        )
+        DS_indicator = self._idx_from_zero(
+            loadtxt(self._file_path("graph_indicator"), delimiter=",").astype(
+                int
+            )
+        )
+
+        if os.path.exists(self._file_path("graph_labels")):
+            DS_graph_labels = self._idx_reset(
+                loadtxt(self._file_path("graph_labels"), delimiter=",").astype(
+                    int
+                )
+            )
+            self.num_labels = int(max(DS_graph_labels) + 1)
+            self.graph_labels = F.tensor(DS_graph_labels)
+        elif os.path.exists(self._file_path("graph_attributes")):
+            DS_graph_labels = loadtxt(
+                self._file_path("graph_attributes"), delimiter=","
+            ).astype(float)
+            self.num_labels = None
+            self.graph_labels = F.tensor(DS_graph_labels)
+        else:
+            raise Exception("Unknown graph label or graph attributes")
+
+        g = dgl_graph(([], []))
+        g.add_nodes(int(DS_edge_list.max()) + 1)
+        g.add_edges(DS_edge_list[:, 0], DS_edge_list[:, 1])
+
+        node_idx_list = []
+        self.max_num_node = 0
+        for idx in range(np.max(DS_indicator) + 1):
+            node_idx = np.where(DS_indicator == idx)
+            node_idx_list.append(node_idx[0])
+            if len(node_idx[0]) > self.max_num_node:
+                self.max_num_node = len(node_idx[0])
+
+        self.attr_dict = {
+            "node_labels": ("ndata", "node_labels"),
+            "node_attributes": ("ndata", "node_attr"),
+            "edge_labels": ("edata", "edge_labels"),
+            "edge_attributes": ("edata", "node_labels"),
+        }
+
+        for filename, field_name in self.attr_dict.items():
+            try:
+                data = loadtxt(self._file_path(filename), delimiter=",")
+                if "label" in filename:
+                    data = F.tensor(self._idx_from_zero(data))
+                else:
+                    data = F.tensor(data)
+                getattr(g, field_name[0])[field_name[1]] = data
+            except IOError:
+                pass
+
+        self.graph_lists = [g.subgraph(node_idx) for node_idx in node_idx_list]
+
+    @property
+    def graph_path(self):
+        return os.path.join(self.save_path, "tu_{}.bin".format(self.name))
+
+    @property
+    def info_path(self):
+        return os.path.join(self.save_path, "tu_{}.pkl".format(self.name))
+
+    def save(self):
+        label_dict = {"labels": self.graph_labels}
+        info_dict = {
+            "max_num_node": self.max_num_node,
+            "num_labels": self.num_labels,
+        }
+        save_graphs(str(self.graph_path), self.graph_lists, label_dict)
+        save_info(str(self.info_path), info_dict)
+
+    def load(self):
+        graphs, label_dict = load_graphs(str(self.graph_path))
+        info_dict = load_info(str(self.info_path))
+
+        self.graph_lists = graphs
+        self.graph_labels = label_dict["labels"]
+        self.max_num_node = info_dict["max_num_node"]
+        self.num_labels = info_dict["num_labels"]
+
+    def has_cache(self):
+        if os.path.exists(self.graph_path) and os.path.exists(self.info_path):
+            return True
+        return False
+
+    def __getitem__(self, idx):
+        """Get the idx-th sample.
+
+        Parameters
+        ---------
+        idx : int
+            The sample index.
+
+        Returns
+        -------
+        (:class:`dgl.DGLGraph`, Tensor)
+            Graph with node feature stored in ``feat`` field and node label in ``node_labels`` if available.
+            And its label.
+        """
+        g = self.graph_lists[idx]
+        if self._transform is not None:
+            g = self._transform(g)
+        return g, self.graph_labels[idx]
+
+    def __len__(self):
+        """Return the number of graphs in the dataset."""
+        return len(self.graph_lists)
+
+    def _file_path(self, category):
+        return os.path.join(
+            self.raw_path, self.name, "{}_{}.txt".format(self.name, category)
+        )
+
+    @staticmethod
+    def _idx_from_zero(idx_tensor):
+        return idx_tensor - np.min(idx_tensor)
+
+    @staticmethod
+    def _idx_reset(idx_tensor):
+        """Maps n unique labels to {0, ..., n-1} in an ordered fashion."""
+        labels = np.unique(idx_tensor)
+        relabel_map = {x: i for i, x in enumerate(labels)}
+        new_idx_tensor = np.vectorize(relabel_map.get)(idx_tensor)
+        return new_idx_tensor
+
+    def statistics(self):
+        return (
+            self.graph_lists[0].ndata["feat"].shape[1],
+            self.num_labels,
+            self.max_num_node,
+        )
+
+    @property
+    def num_classes(self):
+        return self.num_labels

python/dgl/data/utils.py

@@ -0,0 +1,683 @@
+"""Dataset utilities."""
+from __future__ import absolute_import
+
+import errno
+import hashlib
+import os
+import pickle
+import sys
+import warnings
+
+import networkx.algorithms as A
+
+import numpy as np
+import requests
+from tqdm.auto import tqdm
+
+from .. import backend as F
+from .graph_serialize import load_graphs, load_labels, save_graphs
+from .tensor_serialize import load_tensors, save_tensors
+
+__all__ = [
+    "loadtxt",
+    "download",
+    "check_sha1",
+    "extract_archive",
+    "get_download_dir",
+    "Subset",
+    "split_dataset",
+    "save_graphs",
+    "load_graphs",
+    "load_labels",
+    "save_tensors",
+    "load_tensors",
+    "add_nodepred_split",
+    "add_node_property_split",
+    "mask_nodes_by_property",
+]
+
+
+def loadtxt(path, delimiter, dtype=None):
+    try:
+        import pandas as pd
+
+        df = pd.read_csv(path, delimiter=delimiter, header=None)
+        return df.values
+    except ImportError:
+        warnings.warn(
+            "Pandas is not installed, now using numpy.loadtxt to load data, "
+            "which could be extremely slow. Accelerate by installing pandas"
+        )
+        return np.loadtxt(path, delimiter=delimiter)
+
+
+def _get_dgl_url(file_url):
+    """Get DGL online url for download."""
+    dgl_repo_url = "https://data.dgl.ai/"
+    repo_url = os.environ.get("DGL_REPO", dgl_repo_url)
+    if repo_url[-1] != "/":
+        repo_url = repo_url + "/"
+    return repo_url + file_url
+
+
+def split_dataset(dataset, frac_list=None, shuffle=False, random_state=None):
+    """Split dataset into training, validation and test set.
+
+    Parameters
+    ----------
+    dataset
+        We assume ``len(dataset)`` gives the number of datapoints and ``dataset[i]``
+        gives the ith datapoint.
+    frac_list : list or None, optional
+        A list of length 3 containing the fraction to use for training,
+        validation and test. If None, we will use [0.8, 0.1, 0.1].
+    shuffle : bool, optional
+        By default we perform a consecutive split of the dataset. If True,
+        we will first randomly shuffle the dataset.
+    random_state : None, int or array_like, optional
+        Random seed used to initialize the pseudo-random number generator.
+        Can be any integer between 0 and 2**32 - 1 inclusive, an array
+        (or other sequence) of such integers, or None (the default).
+        If seed is None, then RandomState will try to read data from /dev/urandom
+        (or the Windows analogue) if available or seed from the clock otherwise.
+
+    Returns
+    -------
+    list of length 3
+        Subsets for training, validation and test.
+    """
+    from itertools import accumulate
+
+    if frac_list is None:
+        frac_list = [0.8, 0.1, 0.1]
+    frac_list = np.asarray(frac_list)
+    assert np.allclose(
+        np.sum(frac_list), 1.0
+    ), "Expect frac_list sum to 1, got {:.4f}".format(np.sum(frac_list))
+    num_data = len(dataset)
+    lengths = (num_data * frac_list).astype(int)
+    lengths[-1] = num_data - np.sum(lengths[:-1])
+    if shuffle:
+        indices = np.random.RandomState(seed=random_state).permutation(num_data)
+    else:
+        indices = np.arange(num_data)
+    return [
+        Subset(dataset, indices[offset - length : offset])
+        for offset, length in zip(accumulate(lengths), lengths)
+    ]
+
+
+def download(
+    url,
+    path=None,
+    overwrite=True,
+    sha1_hash=None,
+    retries=5,
+    verify_ssl=True,
+    log=True,
+):
+    """Download a given URL.
+
+    Codes borrowed from mxnet/gluon/utils.py
+
+    Parameters
+    ----------
+    url : str
+        URL to download.
+    path : str, optional
+        Destination path to store downloaded file. By default stores to the
+        current directory with the same name as in url.
+    overwrite : bool, optional
+        Whether to overwrite the destination file if it already exists.
+        By default always overwrites the downloaded file.
+    sha1_hash : str, optional
+        Expected sha1 hash in hexadecimal digits. Will ignore existing file when hash is specified
+        but doesn't match.
+    retries : integer, default 5
+        The number of times to attempt downloading in case of failure or non 200 return codes.
+    verify_ssl : bool, default True
+        Verify SSL certificates.
+    log : bool, default True
+        Whether to print the progress for download
+
+    Returns
+    -------
+    str
+        The file path of the downloaded file.
+    """
+    if path is None:
+        fname = url.split("/")[-1]
+        # Empty filenames are invalid
+        assert fname, (
+            "Can't construct file-name from this URL. "
+            "Please set the `path` option manually."
+        )
+    else:
+        path = os.path.expanduser(path)
+        if os.path.isdir(path):
+            fname = os.path.join(path, url.split("/")[-1])
+        else:
+            fname = path
+    assert retries >= 0, "Number of retries should be at least 0"
+
+    if not verify_ssl:
+        warnings.warn(
+            "Unverified HTTPS request is being made (verify_ssl=False). "
+            "Adding certificate verification is strongly advised."
+        )
+
+    if (
+        overwrite
+        or not os.path.exists(fname)
+        or (sha1_hash and not check_sha1(fname, sha1_hash))
+    ):
+        dirname = os.path.dirname(os.path.abspath(os.path.expanduser(fname)))
+        if not os.path.exists(dirname):
+            os.makedirs(dirname)
+        while retries + 1 > 0:
+            # Disable pyling too broad Exception
+            # pylint: disable=W0703
+            try:
+                if log:
+                    print("Downloading %s from %s..." % (fname, url))
+                r = requests.get(url, stream=True, verify=verify_ssl)
+                if r.status_code != 200:
+                    raise RuntimeError("Failed downloading url %s" % url)
+                # Get the total file size.
+                total_size = int(r.headers.get("content-length", 0))
+                with tqdm(
+                    total=total_size, unit="B", unit_scale=True, desc=fname
+                ) as bar:
+                    with open(fname, "wb") as f:
+                        for chunk in r.iter_content(chunk_size=1024):
+                            if chunk:  # filter out keep-alive new chunks
+                                f.write(chunk)
+                                bar.update(len(chunk))
+                if sha1_hash and not check_sha1(fname, sha1_hash):
+                    raise UserWarning(
+                        "File {} is downloaded but the content hash does not match."
+                        " The repo may be outdated or download may be incomplete. "
+                        'If the "repo_url" is overridden, consider switching to '
+                        "the default repo.".format(fname)
+                    )
+                break
+            except Exception as e:
+                retries -= 1
+                if retries <= 0:
+                    raise e
+                else:
+                    if log:
+                        print(
+                            "download failed, retrying, {} attempt{} left".format(
+                                retries, "s" if retries > 1 else ""
+                            )
+                        )
+
+    return fname
+
+
+def check_sha1(filename, sha1_hash):
+    """Check whether the sha1 hash of the file content matches the expected hash.
+
+    Codes borrowed from mxnet/gluon/utils.py
+
+    Parameters
+    ----------
+    filename : str
+        Path to the file.
+    sha1_hash : str
+        Expected sha1 hash in hexadecimal digits.
+
+    Returns
+    -------
+    bool
+        Whether the file content matches the expected hash.
+    """
+    sha1 = hashlib.sha1()
+    with open(filename, "rb") as f:
+        while True:
+            data = f.read(1048576)
+            if not data:
+                break
+            sha1.update(data)
+
+    return sha1.hexdigest() == sha1_hash
+
+
+def extract_archive(file, target_dir, overwrite=True):
+    """Extract archive file.
+
+    Parameters
+    ----------
+    file : str
+        Absolute path of the archive file.
+    target_dir : str
+        Target directory of the archive to be uncompressed.
+    overwrite : bool, default True
+        Whether to overwrite the contents inside the directory.
+        By default always overwrites.
+    """
+    if os.path.exists(target_dir) and not overwrite:
+        return
+    print("Extracting file to {}".format(target_dir))
+    if (
+        file.endswith(".tar.gz")
+        or file.endswith(".tar")
+        or file.endswith(".tgz")
+    ):
+        import tarfile
+
+        with tarfile.open(file, "r") as archive:
+
+            def is_within_directory(directory, target):
+                abs_directory = os.path.abspath(directory)
+                abs_target = os.path.abspath(target)
+                prefix = os.path.commonprefix([abs_directory, abs_target])
+                return prefix == abs_directory
+
+            def safe_extract(
+                tar, path=".", members=None, *, numeric_owner=False
+            ):
+                for member in tar.getmembers():
+                    member_path = os.path.join(path, member.name)
+                    if not is_within_directory(path, member_path):
+                        raise Exception("Attempted Path Traversal in Tar File")
+                tar.extractall(path, members, numeric_owner=numeric_owner)
+
+            safe_extract(archive, path=target_dir)
+    elif file.endswith(".gz"):
+        import gzip
+        import shutil
+
+        with gzip.open(file, "rb") as f_in:
+            target_file = os.path.join(target_dir, os.path.basename(file)[:-3])
+            with open(target_file, "wb") as f_out:
+                shutil.copyfileobj(f_in, f_out)
+    elif file.endswith(".zip"):
+        import zipfile
+
+        with zipfile.ZipFile(file, "r") as archive:
+            archive.extractall(path=target_dir)
+    else:
+        raise Exception("Unrecognized file type: " + file)
+
+
+def get_download_dir():
+    """Get the absolute path to the download directory.
+
+    Returns
+    -------
+    dirname : str
+        Path to the download directory
+    """
+    default_dir = os.path.join(os.path.expanduser("~"), ".dgl")
+    dirname = os.environ.get("DGL_DOWNLOAD_DIR", default_dir)
+    if not os.path.exists(dirname):
+        os.makedirs(dirname)
+    return dirname
+
+
+def makedirs(path):
+    try:
+        os.makedirs(os.path.expanduser(os.path.normpath(path)))
+    except OSError as e:
+        if e.errno != errno.EEXIST and os.path.isdir(path):
+            raise e
+
+
+def save_info(path, info):
+    """Save dataset related information into disk.
+
+    Parameters
+    ----------
+    path : str
+        File to save information.
+    info : dict
+        A python dict storing information to save on disk.
+    """
+    with open(path, "wb") as pf:
+        pickle.dump(info, pf)
+
+
+def load_info(path):
+    """Load dataset related information from disk.
+
+    Parameters
+    ----------
+    path : str
+        File to load information from.
+
+    Returns
+    -------
+    info : dict
+        A python dict storing information loaded from disk.
+    """
+    with open(path, "rb") as pf:
+        info = pickle.load(pf)
+    return info
+
+
+def deprecate_property(old, new):
+    warnings.warn(
+        "Property {} will be deprecated, please use {} instead.".format(
+            old, new
+        )
+    )
+
+
+def deprecate_function(old, new):
+    warnings.warn(
+        "Function {} will be deprecated, please use {} instead.".format(
+            old, new
+        )
+    )
+
+
+def deprecate_class(old, new):
+    warnings.warn(
+        "Class {} will be deprecated, please use {} instead.".format(old, new)
+    )
+
+
+def idx2mask(idx, len):
+    """Create mask."""
+    mask = np.zeros(len)
+    mask[idx] = 1
+    return mask
+
+
+def generate_mask_tensor(mask):
+    """Generate mask tensor according to different backend
+    For torch and tensorflow, it will create a bool tensor
+    For mxnet, it will create a float tensor
+    Parameters
+    ----------
+    mask: numpy ndarray
+        input mask tensor
+    """
+    assert isinstance(mask, np.ndarray), (
+        "input for generate_mask_tensor" "should be an numpy ndarray"
+    )
+    if F.backend_name == "mxnet":
+        return F.tensor(mask, dtype=F.data_type_dict["float32"])
+    else:
+        return F.tensor(mask, dtype=F.data_type_dict["bool"])
+
+
+class Subset(object):
+    """Subset of a dataset at specified indices
+
+    Code adapted from PyTorch.
+
+    Parameters
+    ----------
+    dataset
+        dataset[i] should return the ith datapoint
+    indices : list
+        List of datapoint indices to construct the subset
+    """
+
+    def __init__(self, dataset, indices):
+        self.dataset = dataset
+        self.indices = indices
+
+    def __getitem__(self, item):
+        """Get the datapoint indexed by item
+
+        Returns
+        -------
+        tuple
+            datapoint
+        """
+        return self.dataset[self.indices[item]]
+
+    def __len__(self):
+        """Get subset size
+
+        Returns
+        -------
+        int
+            Number of datapoints in the subset
+        """
+        return len(self.indices)
+
+
+def add_nodepred_split(dataset, ratio, ntype=None):
+    """Split the given dataset into training, validation and test sets for
+    transductive node predction task.
+
+    It adds three node mask arrays ``'train_mask'``, ``'val_mask'`` and ``'test_mask'``,
+    to each graph in the dataset. Each sample in the dataset thus must be a :class:`DGLGraph`.
+
+    Fix the random seed of NumPy to make the result deterministic::
+
+        numpy.random.seed(42)
+
+    Parameters
+    ----------
+    dataset : DGLDataset
+        The dataset to modify.
+    ratio : (float, float, float)
+        Split ratios for training, validation and test sets. Must sum to one.
+    ntype : str, optional
+        The node type to add mask for.
+
+    Examples
+    --------
+    >>> dataset = dgl.data.AmazonCoBuyComputerDataset()
+    >>> print('train_mask' in dataset[0].ndata)
+    False
+    >>> dgl.data.utils.add_nodepred_split(dataset, [0.8, 0.1, 0.1])
+    >>> print('train_mask' in dataset[0].ndata)
+    True
+    """
+    if len(ratio) != 3:
+        raise ValueError(
+            f"Split ratio must be a float triplet but got {ratio}."
+        )
+    for i in range(len(dataset)):
+        g = dataset[i]
+        n = g.num_nodes(ntype)
+        idx = np.arange(0, n)
+        np.random.shuffle(idx)
+        n_train, n_val, n_test = (
+            int(n * ratio[0]),
+            int(n * ratio[1]),
+            int(n * ratio[2]),
+        )
+        train_mask = generate_mask_tensor(idx2mask(idx[:n_train], n))
+        val_mask = generate_mask_tensor(
+            idx2mask(idx[n_train : n_train + n_val], n)
+        )
+        test_mask = generate_mask_tensor(idx2mask(idx[n_train + n_val :], n))
+        g.nodes[ntype].data["train_mask"] = train_mask
+        g.nodes[ntype].data["val_mask"] = val_mask
+        g.nodes[ntype].data["test_mask"] = test_mask
+
+
+def mask_nodes_by_property(property_values, part_ratios, random_seed=None):
+    """Provide the split masks for a node split with distributional shift based on a given
+    node property, as proposed in `Evaluating Robustness and Uncertainty of Graph Models
+    Under Structural Distributional Shifts <https://arxiv.org/abs/2302.13875>`__
+
+    It considers the in-distribution (ID) and out-of-distribution (OOD) subsets of nodes.
+    The ID subset includes training, validation and testing parts, while the OOD subset
+    includes validation and testing parts. It sorts the nodes in the ascending order of
+    their property values, splits them into 5 non-intersecting parts, and creates 5
+    associated node mask arrays:
+        - 3 for the ID nodes: ``'in_train_mask'``, ``'in_valid_mask'``, ``'in_test_mask'``,
+        - and 2 for the OOD nodes: ``'out_valid_mask'``, ``'out_test_mask'``.
+
+    Parameters
+    ----------
+    property_values : numpy ndarray
+        The node property (float) values by which the dataset will be split.
+        The length of the array must be equal to the number of nodes in graph.
+    part_ratios : list
+        A list of 5 ratios for training, ID validation, ID test,
+        OOD validation, OOD testing parts. The values in the list must sum to one.
+    random_seed : int, optional
+        Random seed to fix for the initial permutation of nodes. It is
+        used to create a random order for the nodes that have the same
+        property values or belong to the ID subset. (default: None)
+
+    Returns
+    ----------
+    split_masks : dict
+        A python dict storing the mask names as keys and the corresponding
+        node mask arrays as values.
+
+    Examples
+    --------
+    >>> num_nodes = 1000
+    >>> property_values = np.random.uniform(size=num_nodes)
+    >>> part_ratios = [0.3, 0.1, 0.1, 0.3, 0.2]
+    >>> split_masks = dgl.data.utils.mask_nodes_by_property(property_values, part_ratios)
+    >>> print('in_valid_mask' in split_masks)
+    True
+    """
+
+    num_nodes = len(property_values)
+    part_sizes = np.round(num_nodes * np.array(part_ratios)).astype(int)
+    part_sizes[-1] -= np.sum(part_sizes) - num_nodes
+
+    generator = np.random.RandomState(random_seed)
+    permutation = generator.permutation(num_nodes)
+
+    node_indices = np.arange(num_nodes)[permutation]
+    property_values = property_values[permutation]
+    in_distribution_size = np.sum(part_sizes[:3])
+
+    node_indices_ordered = node_indices[np.argsort(property_values)]
+    node_indices_ordered[:in_distribution_size] = generator.permutation(
+        node_indices_ordered[:in_distribution_size]
+    )
+
+    sections = np.cumsum(part_sizes)
+    node_split = np.split(node_indices_ordered, sections)[:-1]
+    mask_names = [
+        "in_train_mask",
+        "in_valid_mask",
+        "in_test_mask",
+        "out_valid_mask",
+        "out_test_mask",
+    ]
+    split_masks = {}
+
+    for mask_name, node_indices in zip(mask_names, node_split):
+        split_mask = idx2mask(node_indices, num_nodes)
+        split_masks[mask_name] = generate_mask_tensor(split_mask)
+
+    return split_masks
+
+
+def add_node_property_split(
+    dataset, part_ratios, property_name, ascending=True, random_seed=None
+):
+    """Create a node split with distributional shift based on a given node property,
+    as proposed in `Evaluating Robustness and Uncertainty of Graph Models Under
+    Structural Distributional Shifts <https://arxiv.org/abs/2302.13875>`__
+
+    It splits the nodes of each graph in the given dataset into 5 non-intersecting
+    parts based on their structural properties. This can be used for transductive node
+    prediction task with distributional shifts.
+
+    It considers the in-distribution (ID) and out-of-distribution (OOD) subsets of nodes.
+    The ID subset includes training, validation and testing parts, while the OOD subset
+    includes validation and testing parts. As a result, it creates 5 associated node mask
+    arrays for each graph:
+        - 3 for the ID nodes: ``'in_train_mask'``, ``'in_valid_mask'``, ``'in_test_mask'``,
+        - and 2 for the OOD nodes: ``'out_valid_mask'``, ``'out_test_mask'``.
+
+    This function implements 3 particular strategies for inducing distributional shifts
+    in graph — based on **popularity**, **locality** or **density**.
+
+    Parameters
+    ----------
+    dataset : :class:`~DGLDataset` or list of :class:`~dgl.DGLGraph`
+        The dataset to induce structural distributional shift.
+    part_ratios : list
+        A list of 5 ratio values for training, ID validation, ID test,
+        OOD validation and OOD test parts. The values must sum to 1.0.
+    property_name : str
+        The name of the node property to be used, which must be
+        ``'popularity'``, ``'locality'`` or ``'density'``.
+    ascending : bool, optional
+        Whether to sort nodes in the ascending order of the node property,
+        so that nodes with greater values of the property are considered
+        to be OOD (default: True)
+    random_seed : int, optional
+        Random seed to fix for the initial permutation of nodes. It is
+        used to create a random order for the nodes that have the same
+        property values or belong to the ID subset. (default: None)
+
+    Examples
+    --------
+    >>> dataset = dgl.data.AmazonCoBuyComputerDataset()
+    >>> print('in_valid_mask' in dataset[0].ndata)
+    False
+    >>> part_ratios = [0.3, 0.1, 0.1, 0.3, 0.2]
+    >>> property_name = 'popularity'
+    >>> dgl.data.utils.add_node_property_split(dataset, part_ratios, property_name)
+    >>> print('in_valid_mask' in dataset[0].ndata)
+    True
+    """
+
+    assert property_name in [
+        "popularity",
+        "locality",
+        "density",
+    ], "The name of property has to be 'popularity', 'locality', or 'density'"
+
+    assert len(part_ratios) == 5, "part_ratios must contain 5 values"
+
+    import networkx as nx
+
+    for idx in range(len(dataset)):
+        graph_dgl = dataset[idx]
+        graph_nx = nx.Graph(graph_dgl.to_networkx())
+
+        compute_property_fn = _property_name_to_compute_fn[property_name]
+        property_values = compute_property_fn(graph_nx, ascending)
+
+        node_masks = mask_nodes_by_property(
+            property_values, part_ratios, random_seed
+        )
+
+        for mask_name, node_mask in node_masks.items():
+            graph_dgl.ndata[mask_name] = node_mask
+
+
+def _compute_popularity_property(graph_nx, ascending=True):
+    direction = -1 if ascending else 1
+    property_values = direction * np.array(list(A.pagerank(graph_nx).values()))
+    return property_values
+
+
+def _compute_locality_property(graph_nx, ascending=True):
+    num_nodes = graph_nx.number_of_nodes()
+    pagerank_values = np.array(list(A.pagerank(graph_nx).values()))
+
+    personalization = dict(zip(range(num_nodes), [0.0] * num_nodes))
+    personalization[np.argmax(pagerank_values)] = 1.0
+
+    direction = -1 if ascending else 1
+    property_values = direction * np.array(
+        list(A.pagerank(graph_nx, personalization=personalization).values())
+    )
+    return property_values
+
+
+def _compute_density_property(graph_nx, ascending=True):
+    direction = -1 if ascending else 1
+    property_values = direction * np.array(
+        list(A.clustering(graph_nx).values())
+    )
+    return property_values
+
+
+_property_name_to_compute_fn = {
+    "popularity": _compute_popularity_property,
+    "locality": _compute_locality_property,
+    "density": _compute_density_property,
+}

python/dgl/data/wikics.py

@@ -0,0 +1,173 @@
+"""Wiki-CS Dataset"""
+import itertools
+import json
+import os
+
+import numpy as np
+
+from .. import backend as F
+from ..convert import graph
+from ..transforms import reorder_graph, to_bidirected
+from .dgl_dataset import DGLBuiltinDataset
+from .utils import _get_dgl_url, generate_mask_tensor, load_graphs, save_graphs
+
+
+class WikiCSDataset(DGLBuiltinDataset):
+    r"""Wiki-CS is a Wikipedia-based dataset for node classification from `Wiki-CS: A Wikipedia-Based
+    Benchmark for Graph Neural Networks <https://arxiv.org/abs/2007.02901v2>`_
+
+    The dataset consists of nodes corresponding to Computer Science articles, with edges based on
+    hyperlinks and 10 classes representing different branches of the field.
+
+    WikiCS dataset statistics:
+
+    - Nodes: 11,701
+    - Edges: 431,726 (note that the original dataset has 216,123 edges but DGL adds
+      the reverse edges and removes the duplicate edges, hence with a different number)
+    - Number of classes: 10
+    - Node feature size: 300
+    - Number of different train, validation, stopping splits: 20
+    - Number of test split: 1
+
+    Parameters
+    ----------
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset.
+        Default: False
+    verbose : bool
+        Whether to print out progress information.
+        Default: False
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of node classes
+
+    Examples
+    --------
+    >>> from dgl.data import WikiCSDataset
+    >>> dataset = WikiCSDataset()
+    >>> dataset.num_classes
+    10
+    >>> g = dataset[0]
+    >>> # get node feature
+    >>> feat = g.ndata['feat']
+    >>> # get node labels
+    >>> labels = g.ndata['label']
+    >>> # get data split
+    >>> train_mask = g.ndata['train_mask']
+    >>> val_mask = g.ndata['val_mask']
+    >>> stopping_mask = g.ndata['stopping_mask']
+    >>> test_mask = g.ndata['test_mask']
+    >>> # The shape of train, val and stopping masks are (num_nodes, num_splits).
+    >>> # The num_splits is the number of different train, validation, stopping splits.
+    >>> # Due to the number of test spilt is 1, the shape of test mask is (num_nodes,).
+    >>> print(train_mask.shape, val_mask.shape, stopping_mask.shape)
+    (11701, 20) (11701, 20) (11701, 20)
+    >>> print(test_mask.shape)
+    (11701,)
+    """
+
+    def __init__(
+        self, raw_dir=None, force_reload=False, verbose=False, transform=None
+    ):
+        _url = _get_dgl_url("dataset/wiki_cs.zip")
+        super(WikiCSDataset, self).__init__(
+            name="wiki_cs",
+            raw_dir=raw_dir,
+            url=_url,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        """process raw data to graph, labels and masks"""
+        with open(os.path.join(self.raw_path, "data.json")) as f:
+            data = json.load(f)
+        features = F.tensor(np.array(data["features"]), dtype=F.float32)
+        labels = F.tensor(np.array(data["labels"]), dtype=F.int64)
+
+        train_masks = np.array(data["train_masks"], dtype=bool).T
+        val_masks = np.array(data["val_masks"], dtype=bool).T
+        stopping_masks = np.array(data["stopping_masks"], dtype=bool).T
+        test_mask = np.array(data["test_mask"], dtype=bool)
+
+        edges = [[(i, j) for j in js] for i, js in enumerate(data["links"])]
+        edges = np.array(list(itertools.chain(*edges)))
+        src, dst = edges[:, 0], edges[:, 1]
+
+        g = graph((src, dst))
+        g = to_bidirected(g)
+
+        g.ndata["feat"] = features
+        g.ndata["label"] = labels
+        g.ndata["train_mask"] = generate_mask_tensor(train_masks)
+        g.ndata["val_mask"] = generate_mask_tensor(val_masks)
+        g.ndata["stopping_mask"] = generate_mask_tensor(stopping_masks)
+        g.ndata["test_mask"] = generate_mask_tensor(test_mask)
+
+        g = reorder_graph(
+            g,
+            node_permute_algo="rcmk",
+            edge_permute_algo="dst",
+            store_ids=False,
+        )
+
+        self._graph = g
+
+    def has_cache(self):
+        graph_path = os.path.join(self.save_path, "dgl_graph.bin")
+        return os.path.exists(graph_path)
+
+    def save(self):
+        graph_path = os.path.join(self.save_path, "dgl_graph.bin")
+        save_graphs(graph_path, self._graph)
+
+    def load(self):
+        graph_path = os.path.join(self.save_path, "dgl_graph.bin")
+        g, _ = load_graphs(graph_path)
+        self._graph = g[0]
+
+    @property
+    def num_classes(self):
+        return 10
+
+    def __len__(self):
+        r"""The number of graphs in the dataset."""
+        return 1
+
+    def __getitem__(self, idx):
+        r"""Get graph object
+
+        Parameters
+        ----------
+        idx : int
+            Item index, WikiCSDataset has only one graph object
+
+        Returns
+        -------
+        :class:`dgl.DGLGraph`
+
+            The graph contains:
+
+            - ``ndata['feat']``: node features
+            - ``ndata['label']``: node labels
+            - ``ndata['train_mask']``: train mask is for retrieving the nodes for training.
+            - ``ndata['val_mask']``: val mask is for retrieving the nodes for hyperparameter tuning.
+            - ``ndata['stopping_mask']``: stopping mask is for retrieving the nodes for early stopping criterion.
+            - ``ndata['test_mask']``: test mask is for retrieving the nodes for testing.
+
+        """
+        assert idx == 0, "This dataset has only one graph"
+        if self._transform is None:
+            return self._graph
+        else:
+            return self._transform(self._graph)

python/dgl/data/yelp.py

@@ -0,0 +1,177 @@
+"""Yelp Dataset"""
+import json
+import os
+
+import numpy as np
+import scipy.sparse as sp
+
+from .. import backend as F
+from ..convert import from_scipy
+from ..transforms import reorder_graph
+from .dgl_dataset import DGLBuiltinDataset
+from .utils import _get_dgl_url, generate_mask_tensor, load_graphs, save_graphs
+
+
+class YelpDataset(DGLBuiltinDataset):
+    r"""Yelp dataset for node classification from `GraphSAINT: Graph Sampling Based Inductive
+    Learning Method <https://arxiv.org/abs/1907.04931>`_
+
+    The task of this dataset is categorizing types of businesses based on customer reviewers and
+    friendship.
+
+    Yelp dataset statistics:
+
+    - Nodes: 716,847
+    - Edges: 13,954,819
+    - Number of classes: 100 (Multi-class)
+    - Node feature size: 300
+
+    Parameters
+    ----------
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset.
+        Default: False
+    verbose : bool
+        Whether to print out progress information.
+        Default: False
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+    reorder : bool
+        Whether to reorder the graph using :func:`~dgl.reorder_graph`.
+        Default: False.
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of node classes
+
+    Examples
+    --------
+    >>> dataset = YelpDataset()
+    >>> dataset.num_classes
+    100
+    >>> g = dataset[0]
+    >>> # get node feature
+    >>> feat = g.ndata['feat']
+    >>> # get node labels
+    >>> labels = g.ndata['label']
+    >>> # get data split
+    >>> train_mask = g.ndata['train_mask']
+    >>> val_mask = g.ndata['val_mask']
+    >>> test_mask = g.ndata['test_mask']
+    """
+
+    def __init__(
+        self,
+        raw_dir=None,
+        force_reload=False,
+        verbose=False,
+        transform=None,
+        reorder=False,
+    ):
+        _url = _get_dgl_url("dataset/yelp.zip")
+        self._reorder = reorder
+        super(YelpDataset, self).__init__(
+            name="yelp",
+            raw_dir=raw_dir,
+            url=_url,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        """process raw data to graph, labels and masks"""
+        coo_adj = sp.load_npz(os.path.join(self.raw_path, "adj_full.npz"))
+        g = from_scipy(coo_adj)
+
+        features = np.load(os.path.join(self.raw_path, "feats.npy"))
+        features = F.tensor(features, dtype=F.float32)
+
+        y = [-1] * features.shape[0]
+        with open(os.path.join(self.raw_path, "class_map.json")) as f:
+            class_map = json.load(f)
+            for key, item in class_map.items():
+                y[int(key)] = item
+        labels = F.tensor(np.array(y), dtype=F.int64)
+
+        with open(os.path.join(self.raw_path, "role.json")) as f:
+            role = json.load(f)
+
+        train_mask = np.zeros(features.shape[0], dtype=bool)
+        train_mask[role["tr"]] = True
+
+        val_mask = np.zeros(features.shape[0], dtype=bool)
+        val_mask[role["va"]] = True
+
+        test_mask = np.zeros(features.shape[0], dtype=bool)
+        test_mask[role["te"]] = True
+
+        g.ndata["feat"] = features
+        g.ndata["label"] = labels
+        g.ndata["train_mask"] = generate_mask_tensor(train_mask)
+        g.ndata["val_mask"] = generate_mask_tensor(val_mask)
+        g.ndata["test_mask"] = generate_mask_tensor(test_mask)
+
+        if self._reorder:
+            self._graph = reorder_graph(
+                g,
+                node_permute_algo="rcmk",
+                edge_permute_algo="dst",
+                store_ids=False,
+            )
+        else:
+            self._graph = g
+
+    def has_cache(self):
+        graph_path = os.path.join(self.save_path, "dgl_graph.bin")
+        return os.path.exists(graph_path)
+
+    def save(self):
+        graph_path = os.path.join(self.save_path, "dgl_graph.bin")
+        save_graphs(graph_path, self._graph)
+
+    def load(self):
+        graph_path = os.path.join(self.save_path, "dgl_graph.bin")
+        g, _ = load_graphs(graph_path)
+        self._graph = g[0]
+
+    @property
+    def num_classes(self):
+        return 100
+
+    def __len__(self):
+        r"""The number of graphs in the dataset."""
+        return 1
+
+    def __getitem__(self, idx):
+        r"""Get graph object
+
+        Parameters
+        ----------
+        idx : int
+            Item index, FlickrDataset has only one graph object
+
+        Returns
+        -------
+        :class:`dgl.DGLGraph`
+
+            The graph contains:
+
+            - ``ndata['label']``: node label
+            - ``ndata['feat']``: node feature
+            - ``ndata['train_mask']``: mask for training node set
+            - ``ndata['val_mask']``: mask for validation node set
+            - ``ndata['test_mask']``: mask for test node set
+
+        """
+        assert idx == 0, "This dataset has only one graph"
+        if self._transform is None:
+            return self._graph
+        else:
+            return self._transform(self._graph)

python/dgl/data/zinc.py

@@ -0,0 +1,137 @@
+import os
+
+from .dgl_dataset import DGLBuiltinDataset
+from .utils import _get_dgl_url, load_graphs
+
+
+class ZINCDataset(DGLBuiltinDataset):
+    r"""ZINC dataset for the graph regression task.
+
+    A subset (12K) of ZINC molecular graphs (250K) dataset is used to
+    regress a molecular property known as the constrained solubility.
+    For each molecular graph, the node features are the types of heavy
+    atoms, between which the edge features are the types of bonds.
+    Each graph contains 9-37 nodes and 16-84 edges.
+
+    Reference `<https://arxiv.org/pdf/2003.00982.pdf>`_
+
+    Statistics:
+
+    Train examples: 10,000
+    Valid examples: 1,000
+    Test examples: 1,000
+    Average number of nodes: 23.16
+    Average number of edges: 39.83
+    Number of atom types: 28
+    Number of bond types: 4
+
+    Parameters
+    ----------
+    mode : str, optional
+        Should be chosen from ["train", "valid", "test"]
+        Default: "train".
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: "~/.dgl/".
+    force_reload : bool
+        Whether to reload the dataset.
+        Default: False.
+    verbose : bool
+        Whether to print out progress information.
+        Default: False.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    num_atom_types : int
+        Number of atom types.
+    num_bond_types : int
+        Number of bond types.
+
+    Examples
+    ---------
+    >>> from dgl.data import ZINCDataset
+
+    >>> training_set = ZINCDataset(mode="train")
+    >>> training_set.num_atom_types
+    28
+    >>> len(training_set)
+    10000
+    >>> graph, label = training_set[0]
+    >>> graph
+    Graph(num_nodes=29, num_edges=64,
+        ndata_schemes={'feat': Scheme(shape=(), dtype=torch.int64)}
+        edata_schemes={'feat': Scheme(shape=(), dtype=torch.int64)})
+    """
+
+    def __init__(
+        self,
+        mode="train",
+        raw_dir=None,
+        force_reload=False,
+        verbose=False,
+        transform=None,
+    ):
+        self._url = _get_dgl_url("dataset/ZINC12k.zip")
+        self.mode = mode
+
+        super(ZINCDataset, self).__init__(
+            name="zinc",
+            url=self._url,
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        self.load()
+
+    @property
+    def graph_path(self):
+        return os.path.join(self.save_path, "ZincDGL_{}.bin".format(self.mode))
+
+    def has_cache(self):
+        return os.path.exists(self.graph_path)
+
+    def load(self):
+        self._graphs, self._labels = load_graphs(self.graph_path)
+
+    @property
+    def num_atom_types(self):
+        return 28
+
+    @property
+    def num_bond_types(self):
+        return 4
+
+    def __len__(self):
+        return len(self._graphs)
+
+    def __getitem__(self, idx):
+        r"""Get one example by index.
+
+        Parameters
+        ----------
+        idx : int
+            The sample index.
+
+        Returns
+        -------
+        dgl.DGLGraph
+            Each graph contains:
+
+            - ``ndata['feat']``: Types of heavy atoms as node features
+            - ``edata['feat']``: Types of bonds as edge features
+
+        Tensor
+            Constrained solubility as graph label
+        """
+        labels = self._labels["g_label"]
+        if self._transform is None:
+            return self._graphs[idx], labels[idx]
+        else:
+            return self._transform(self._graphs[idx]), labels[idx]