wehub-resource-sync
21 lines
540 B
Python
21 lines
540 B
Python
from dgl.data import CiteseerGraphDataset, CoraGraphDataset, 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
|