"vscode:/vscode.git/clone" did not exist on "344342abdf41c97ddcf3ace2d562e2ff0a5af1f3"
load_dataset.py 1.41 KB
Newer Older
1
2
3
""" load dataset from ogb """

import argparse
4

5
6
from ogb.linkproppred import DglLinkPropPredDataset
from ogb.nodeproppred import DglNodePropPredDataset
7

8
9
import dgl

10
11
12

def load_from_ogbl_with_name(name):
    choices = ["ogbl-collab", "ogbl-ddi", "ogbl-ppa", "ogbl-citation"]
13
14
15
16
    assert name in choices, "name must be selected from " + str(choices)
    dataset = DglLinkPropPredDataset(name)
    return dataset[0]

17
18
19
20
21
22
23
24

def load_from_ogbn_with_name(name):
    choices = [
        "ogbn-products",
        "ogbn-proteins",
        "ogbn-arxiv",
        "ogbn-papers100M",
    ]
25
26
27
28
    assert name in choices, "name must be selected from " + str(choices)
    dataset, label = DglNodePropPredDataset(name)[0]
    return dataset

29

30
if __name__ == "__main__":
31
    """load datasets as net.txt format"""
32
    parser = argparse.ArgumentParser()
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
    parser.add_argument(
        "--name",
        type=str,
        choices=[
            "ogbl-collab",
            "ogbl-ddi",
            "ogbl-ppa",
            "ogbl-citation",
            "ogbn-products",
            "ogbn-proteins",
            "ogbn-arxiv",
            "ogbn-papers100M",
        ],
        default="ogbl-collab",
        help="name of datasets by ogb",
    )
49
50
51
52
53
54
55
56
    args = parser.parse_args()

    name = args.name
    if name.startswith("ogbl"):
        g = load_from_ogbl_with_name(name=name)
    else:
        g = load_from_ogbn_with_name(name=name)

57
    dgl.save_graphs(name + "-graph.bin", g)