def _test(in_readonly, out_readonly):
elist = [(0, 0), (0, 1), (1, 0), (1, 1), (2, 1), (2, 2)]
num_edges = 7
g = dgl.DGLGraph(elist, readonly=in_readonly)
elist.append((1, 2))
elist = set(elist)
big = dgl.to_bidirected_stale(g, out_readonly)
assert big.number_of_edges() == num_edges
src, dst = big.edges()
eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst))))
assert eset == set(elist)
if args.dataset == 'reddit':
g, _ = load_reddit()
elif args.dataset == 'ogb-product':
g, _ = load_ogb('ogbn-products')
elif args.dataset == 'ogb-paper100M':
g, _ = load_ogb('ogbn-papers100M')
print('load {} takes {:.3f} seconds'.format(args.dataset,
time.time() - start))
print('|V|={}, |E|={}'.format(g.number_of_nodes(), g.number_of_edges()))
print('train: {}, valid: {}, test: {}'.format(
th.sum(g.ndata['train_mask']), th.sum(g.ndata['val_mask']),
th.sum(g.ndata['test_mask'])))
if args.balance_train:
balance_ntypes = g.ndata['train_mask']
else:
balance_ntypes = None
if args.undirected:
sym_g = dgl.to_bidirected_stale(g, readonly=True)
for key in g.ndata:
sym_g.ndata[key] = g.ndata[key]
g = sym_g
dgl.distributed.partition_graph(g,
args.dataset,
args.num_parts,
'data',
part_method=args.part_method,
balance_ntypes=balance_ntypes,
balance_edges=args.balance_edges)