chore: import upstream snapshot with attribution

2026-07-13 13:35:51 +08:00
commit c36a561cd8
2172 changed files with 455595 additions and 0 deletions
@@ -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".')