dmlc--dgl
Hao Xiong
9488bace20
* line * two lines * update readme * readme * update readme * update Co-authored-by: Zihao Ye <expye@outlook.com> Co-authored-by: xiang song(charlie.song) <classicxsong@gmail.com>
50 行
1.5 KiB
Python
50 行
1.5 KiB
Python
""" load dataset from ogb """
|
|
|
|
import argparse
|
|
from ogb.linkproppred import DglLinkPropPredDataset
|
|
import time
|
|
|
|
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]
|
|
|
|
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()
|
|
|
|
print("loading graph... it might take some time")
|
|
name = args.name
|
|
g = load_from_ogbl_with_name(name=name)
|
|
|
|
try:
|
|
w = g.edata['edge_weight']
|
|
weighted = True
|
|
except:
|
|
weighted = False
|
|
|
|
|
|
edge_num = g.edges()[0].shape[0]
|
|
src = list(g.edges()[0])
|
|
tgt = list(g.edges()[1])
|
|
if weighted:
|
|
weight = list(g.edata['edge_weight'])
|
|
|
|
print("writing...")
|
|
start_time = time.time()
|
|
with open(name + "-net.txt", "w") as f:
|
|
for i in range(edge_num):
|
|
if weighted:
|
|
f.write(str(src[i].item()) + " "\
|
|
+str(tgt[i].item()) + " "\
|
|
+str(weight[i].item()) + "\n")
|
|
else:
|
|
f.write(str(src[i].item()) + " "\
|
|
+str(tgt[i].item()) + " "\
|
|
+"1\n")
|
|
print("writing used time: %d s" % int(time.time() - start_time)) |