dmlc--dgl
Hao Xiong
eef4c05929
* line * two lines * update readme * readme * update readme * update * Implement LINE * readme * readme * typos * update readme Co-authored-by: Zihao Ye <expye@outlook.com> Co-authored-by: xiang song(charlie.song) <classicxsong@gmail.com> Co-authored-by: Jinjing Zhou <VoVAllen@users.noreply.github.com>
36 行
1.3 KiB
Python
36 行
1.3 KiB
Python
""" load dataset from ogb """
|
|
|
|
import argparse
|
|
from ogb.linkproppred import DglLinkPropPredDataset
|
|
from ogb.nodeproppred import DglNodePropPredDataset
|
|
import dgl
|
|
|
|
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]
|
|
|
|
def load_from_ogbn_with_name(name):
|
|
choices = ['ogbn-products', 'ogbn-proteins', 'ogbn-arxiv', 'ogbn-papers100M']
|
|
assert name in choices, "name must be selected from " + str(choices)
|
|
dataset, label = DglNodePropPredDataset(name)[0]
|
|
return dataset
|
|
|
|
if __name__ == "__main__":
|
|
""" load datasets as net.txt format """
|
|
parser = argparse.ArgumentParser()
|
|
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")
|
|
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)
|
|
|
|
dgl.save_graphs(name + "-graph.bin", g) |