"torchvision/vscode:/vscode.git/clone" did not exist on "fbd69f1052292336f78b69804a68ade5a6b2f3b4"
load_dataset.py 1.23 KB
Newer Older
1
2
3
""" load dataset from ogb """

import argparse
4
from ogb.linkproppred import DglLinkPropPredDataset
5

6
7
8
9
10
def load_from_ogbl_with_name(name):    
    choices = ['ogbl-collab', 'ogbl-ddi', 'ogbl-ppa', 'ogbl-citation']
    assert name in choices, "name must be selected from " + str(choices)
    dataset = DglLinkPropPredDataset(name)
    return dataset[0]
11

12
13
14
15
16
17
18
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--name', type=str,
        choices=['ogbl-collab', 'ogbl-ddi', 'ogbl-ppa', 'ogbl-citation'],
        default='ogbl-collab',
        help="name of datasets by ogb")
    args = parser.parse_args()
19

20
21
    name = args.name
    g = load_from_ogbl_with_name(name=name)
22

23
24
25
26
27
    try:
        w = g.edata['edge_weight']
        weighted = True
    except:
        weighted = False
28

29
30
31
32
33
34
35
36
37
38
    with open(name + "-net.txt", "w") as f:
        for i in range(g.edges()[0].shape[0]):
            if weighted:
                f.write(str(g.edges()[0][i].item()) + " "\
                    +str(g.edges()[1][i].item()) + " "\
                    +str(g.edata['edge_weight'][i]) + "\n")
            else:
                f.write(str(g.edges()[0][i].item()) + " "\
                    +str(g.edges()[1][i].item()) + " "\
                    +"1\n")