load_graph.py 1.82 KB
Newer Older
1
2
import torch as th

3
4
5
import dgl


6
def load_reddit(self_loop=True):
7
8
9
    from dgl.data import RedditDataset

    # load reddit data
10
    data = RedditDataset(self_loop=self_loop)
Xiangkun Hu's avatar
Xiangkun Hu committed
11
    g = data[0]
12
13
    g.ndata["features"] = g.ndata.pop("feat")
    g.ndata["labels"] = g.ndata.pop("label")
14
    return g, data.num_classes
15

16
17

def load_ogb(name, root="dataset"):
18
19
    from ogb.nodeproppred import DglNodePropPredDataset

20
    print("load", name)
21
    data = DglNodePropPredDataset(name=name, root=root)
22
    print("finish loading", name)
23
24
25
26
    splitted_idx = data.get_idx_split()
    graph, labels = data[0]
    labels = labels[:, 0]

27
28
29
    graph.ndata["features"] = graph.ndata.pop("feat")
    graph.ndata["labels"] = labels
    in_feats = graph.ndata["features"].shape[1]
Da Zheng's avatar
Da Zheng committed
30
    num_labels = len(th.unique(labels[th.logical_not(th.isnan(labels))]))
31
32

    # Find the node IDs in the training, validation, and test set.
33
34
35
36
37
    train_nid, val_nid, test_nid = (
        splitted_idx["train"],
        splitted_idx["valid"],
        splitted_idx["test"],
    )
38
39
40
41
42
43
    train_mask = th.zeros((graph.number_of_nodes(),), dtype=th.bool)
    train_mask[train_nid] = True
    val_mask = th.zeros((graph.number_of_nodes(),), dtype=th.bool)
    val_mask[val_nid] = True
    test_mask = th.zeros((graph.number_of_nodes(),), dtype=th.bool)
    test_mask[test_nid] = True
44
45
46
47
    graph.ndata["train_mask"] = train_mask
    graph.ndata["val_mask"] = val_mask
    graph.ndata["test_mask"] = test_mask
    print("finish constructing", name)
Da Zheng's avatar
Da Zheng committed
48
    return graph, num_labels
49

50

51
52
53
def inductive_split(g):
    """Split the graph into training graph, validation graph, and test graph by training
    and validation masks.  Suitable for inductive models."""
54
55
    train_g = g.subgraph(g.ndata["train_mask"])
    val_g = g.subgraph(g.ndata["train_mask"] | g.ndata["val_mask"])
56
57
    test_g = g
    return train_g, val_g, test_g