dmlc--dgl
Hengrui Zhang
b2b531e041
* [Example] add implementation of grace * [Doc] add grace in the index file * fix * fix typos Co-authored-by: Mufei Li <mufeili1996@gmail.com>
19 行
538 B
Python
19 行
538 B
Python
from dgl.data import CoraGraphDataset, CiteseerGraphDataset, PubmedGraphDataset
|
|
|
|
def load(name):
|
|
if name == 'cora':
|
|
dataset = CoraGraphDataset()
|
|
elif name == 'citeseer':
|
|
dataset = CiteseerGraphDataset()
|
|
elif name == 'pubmed':
|
|
dataset = PubmedGraphDataset()
|
|
|
|
graph = dataset[0]
|
|
|
|
train_mask = graph.ndata.pop('train_mask')
|
|
test_mask = graph.ndata.pop('test_mask')
|
|
|
|
feat = graph.ndata.pop('feat')
|
|
labels = graph.ndata.pop('label')
|
|
|
|
return graph, feat, labels, train_mask, test_mask |