wehub-resource-sync
421 lines
14 KiB
Python
421 lines
14 KiB
Python
"""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
|