dmlc--dgl
Da Zheng
454 行
14 KiB
Python
454 行
14 KiB
Python
import os
|
|
os.environ['DGLBACKEND'] = 'mxnet'
|
|
import mxnet as mx
|
|
import numpy as np
|
|
from dgl.graph import DGLGraph
|
|
import dgl
|
|
import scipy.sparse as spsp
|
|
|
|
D = 5
|
|
reduce_msg_shapes = set()
|
|
|
|
def check_eq(a, b):
|
|
assert a.shape == b.shape
|
|
assert mx.nd.sum(a == b).asnumpy() == int(np.prod(list(a.shape)))
|
|
|
|
def message_func(edges):
|
|
assert len(edges.src['h'].shape) == 2
|
|
assert edges.src['h'].shape[1] == D
|
|
return {'m' : edges.src['h']}
|
|
|
|
def reduce_func(nodes):
|
|
msgs = nodes.mailbox['m']
|
|
reduce_msg_shapes.add(tuple(msgs.shape))
|
|
assert len(msgs.shape) == 3
|
|
assert msgs.shape[2] == D
|
|
return {'m' : mx.nd.sum(msgs, 1)}
|
|
|
|
def apply_node_func(nodes):
|
|
return {'h' : nodes.data['h'] + nodes.data['m']}
|
|
|
|
def generate_graph(grad=False, readonly=False):
|
|
if readonly:
|
|
row_idx = []
|
|
col_idx = []
|
|
for i in range(1, 9):
|
|
row_idx.append(0)
|
|
col_idx.append(i)
|
|
row_idx.append(i)
|
|
col_idx.append(9)
|
|
row_idx.append(9)
|
|
col_idx.append(0)
|
|
ones = np.ones(shape=(len(row_idx)))
|
|
csr = spsp.csr_matrix((ones, (row_idx, col_idx)), shape=(10, 10))
|
|
g = DGLGraph(csr, readonly=True)
|
|
ncol = mx.nd.random.normal(shape=(10, D))
|
|
ecol = mx.nd.random.normal(shape=(17, D))
|
|
if grad:
|
|
ncol.attach_grad()
|
|
ecol.attach_grad()
|
|
g.ndata['h'] = ncol
|
|
g.edata['w'] = ecol
|
|
g.set_n_initializer(dgl.init.zero_initializer)
|
|
g.set_e_initializer(dgl.init.zero_initializer)
|
|
return g
|
|
else:
|
|
g = DGLGraph()
|
|
g.add_nodes(10) # 10 nodes.
|
|
# create a graph where 0 is the source and 9 is the sink
|
|
for i in range(1, 9):
|
|
g.add_edge(0, i)
|
|
g.add_edge(i, 9)
|
|
# add a back flow from 9 to 0
|
|
g.add_edge(9, 0)
|
|
ncol = mx.nd.random.normal(shape=(10, D))
|
|
ecol = mx.nd.random.normal(shape=(17, D))
|
|
if grad:
|
|
ncol.attach_grad()
|
|
ecol.attach_grad()
|
|
g.ndata['h'] = ncol
|
|
g.edata['w'] = ecol
|
|
g.set_n_initializer(dgl.init.zero_initializer)
|
|
g.set_e_initializer(dgl.init.zero_initializer)
|
|
return g
|
|
|
|
def test_batch_setter_getter():
|
|
def _pfc(x):
|
|
return list(x.asnumpy()[:,0])
|
|
g = generate_graph()
|
|
# set all nodes
|
|
g.set_n_repr({'h' : mx.nd.zeros((10, D))})
|
|
assert _pfc(g.ndata['h']) == [0.] * 10
|
|
# pop nodes
|
|
assert _pfc(g.pop_n_repr('h')) == [0.] * 10
|
|
assert len(g.ndata) == 0
|
|
g.set_n_repr({'h' : mx.nd.zeros((10, D))})
|
|
# set partial nodes
|
|
u = mx.nd.array([1, 3, 5], dtype='int64')
|
|
g.set_n_repr({'h' : mx.nd.ones((3, D))}, u)
|
|
assert _pfc(g.ndata['h']) == [0., 1., 0., 1., 0., 1., 0., 0., 0., 0.]
|
|
# get partial nodes
|
|
u = mx.nd.array([1, 2, 3], dtype='int64')
|
|
assert _pfc(g.get_n_repr(u)['h']) == [1., 0., 1.]
|
|
|
|
'''
|
|
s, d, eid
|
|
0, 1, 0
|
|
1, 9, 1
|
|
0, 2, 2
|
|
2, 9, 3
|
|
0, 3, 4
|
|
3, 9, 5
|
|
0, 4, 6
|
|
4, 9, 7
|
|
0, 5, 8
|
|
5, 9, 9
|
|
0, 6, 10
|
|
6, 9, 11
|
|
0, 7, 12
|
|
7, 9, 13
|
|
0, 8, 14
|
|
8, 9, 15
|
|
9, 0, 16
|
|
'''
|
|
# set all edges
|
|
g.edata['l'] = mx.nd.zeros((17, D))
|
|
assert _pfc(g.edata['l']) == [0.] * 17
|
|
# pop edges
|
|
old_len = len(g.edata)
|
|
assert _pfc(g.pop_e_repr('l')) == [0.] * 17
|
|
assert len(g.edata) == old_len - 1
|
|
g.edata['l'] = mx.nd.zeros((17, D))
|
|
# set partial edges (many-many)
|
|
u = mx.nd.array([0, 0, 2, 5, 9], dtype='int64')
|
|
v = mx.nd.array([1, 3, 9, 9, 0], dtype='int64')
|
|
g.edges[u, v].data['l'] = mx.nd.ones((5, D))
|
|
truth = [0.] * 17
|
|
truth[0] = truth[4] = truth[3] = truth[9] = truth[16] = 1.
|
|
assert _pfc(g.edata['l']) == truth
|
|
# set partial edges (many-one)
|
|
u = mx.nd.array([3, 4, 6], dtype='int64')
|
|
v = mx.nd.array([9], dtype='int64')
|
|
g.edges[u, v].data['l'] = mx.nd.ones((3, D))
|
|
truth[5] = truth[7] = truth[11] = 1.
|
|
assert _pfc(g.edata['l']) == truth
|
|
# set partial edges (one-many)
|
|
u = mx.nd.array([0], dtype='int64')
|
|
v = mx.nd.array([4, 5, 6], dtype='int64')
|
|
g.edges[u, v].data['l'] = mx.nd.ones((3, D))
|
|
truth[6] = truth[8] = truth[10] = 1.
|
|
assert _pfc(g.edata['l']) == truth
|
|
# get partial edges (many-many)
|
|
u = mx.nd.array([0, 6, 0], dtype='int64')
|
|
v = mx.nd.array([6, 9, 7], dtype='int64')
|
|
assert _pfc(g.edges[u, v].data['l']) == [1., 1., 0.]
|
|
# get partial edges (many-one)
|
|
u = mx.nd.array([5, 6, 7], dtype='int64')
|
|
v = mx.nd.array([9], dtype='int64')
|
|
assert _pfc(g.edges[u, v].data['l']) == [1., 1., 0.]
|
|
# get partial edges (one-many)
|
|
u = mx.nd.array([0], dtype='int64')
|
|
v = mx.nd.array([3, 4, 5], dtype='int64')
|
|
assert _pfc(g.edges[u, v].data['l']) == [1., 1., 1.]
|
|
|
|
def test_batch_setter_autograd():
|
|
with mx.autograd.record():
|
|
g = generate_graph(grad=True, readonly=True)
|
|
h1 = g.ndata['h']
|
|
h1.attach_grad()
|
|
# partial set
|
|
v = mx.nd.array([1, 2, 8], dtype='int64')
|
|
hh = mx.nd.zeros((len(v), D))
|
|
hh.attach_grad()
|
|
g.set_n_repr({'h' : hh}, v)
|
|
h2 = g.ndata['h']
|
|
h2.backward(mx.nd.ones((10, D)) * 2)
|
|
check_eq(h1.grad[:,0], mx.nd.array([2., 0., 0., 2., 2., 2., 2., 2., 0., 2.]))
|
|
check_eq(hh.grad[:,0], mx.nd.array([2., 2., 2.]))
|
|
|
|
def test_batch_send():
|
|
g = generate_graph()
|
|
def _fmsg(edges):
|
|
assert edges.src['h'].shape == (5, D)
|
|
return {'m' : edges.src['h']}
|
|
g.register_message_func(_fmsg)
|
|
# many-many send
|
|
u = mx.nd.array([0, 0, 0, 0, 0], dtype='int64')
|
|
v = mx.nd.array([1, 2, 3, 4, 5], dtype='int64')
|
|
g.send((u, v))
|
|
# one-many send
|
|
u = mx.nd.array([0], dtype='int64')
|
|
v = mx.nd.array([1, 2, 3, 4, 5], dtype='int64')
|
|
g.send((u, v))
|
|
# many-one send
|
|
u = mx.nd.array([1, 2, 3, 4, 5], dtype='int64')
|
|
v = mx.nd.array([9], dtype='int64')
|
|
g.send((u, v))
|
|
|
|
def check_batch_recv(readonly):
|
|
# basic recv test
|
|
g = generate_graph(readonly=readonly)
|
|
g.register_message_func(message_func)
|
|
g.register_reduce_func(reduce_func)
|
|
g.register_apply_node_func(apply_node_func)
|
|
u = mx.nd.array([0, 0, 0, 4, 5, 6], dtype='int64')
|
|
v = mx.nd.array([1, 2, 3, 9, 9, 9], dtype='int64')
|
|
reduce_msg_shapes.clear()
|
|
g.send((u, v))
|
|
#g.recv(th.unique(v))
|
|
#assert(reduce_msg_shapes == {(1, 3, D), (3, 1, D)})
|
|
#reduce_msg_shapes.clear()
|
|
|
|
def test_batch_recv():
|
|
check_batch_recv(True)
|
|
check_batch_recv(False)
|
|
|
|
def test_apply_nodes():
|
|
def _upd(nodes):
|
|
return {'h' : nodes.data['h'] * 2}
|
|
g = generate_graph()
|
|
g.register_apply_node_func(_upd)
|
|
old = g.ndata['h']
|
|
g.apply_nodes()
|
|
assert np.allclose((old * 2).asnumpy(), g.ndata['h'].asnumpy())
|
|
u = mx.nd.array([0, 3, 4, 6], dtype=np.int64)
|
|
g.apply_nodes(lambda nodes : {'h' : nodes.data['h'] * 0.}, u)
|
|
h = g.ndata['h'][u].asnumpy()
|
|
assert np.allclose(h, np.zeros(shape=(4, D), dtype=h.dtype))
|
|
|
|
def test_apply_edges():
|
|
def _upd(edges):
|
|
return {'w' : edges.data['w'] * 2}
|
|
g = generate_graph()
|
|
g.register_apply_edge_func(_upd)
|
|
old = g.edata['w']
|
|
g.apply_edges()
|
|
assert np.allclose((old * 2).asnumpy(), g.edata['w'].asnumpy())
|
|
u = mx.nd.array([0, 0, 0, 4, 5, 6], dtype=np.int64)
|
|
v = mx.nd.array([1, 2, 3, 9, 9, 9], dtype=np.int64)
|
|
g.apply_edges(lambda edges : {'w' : edges.data['w'] * 0.}, (u, v))
|
|
eid = g.edge_ids(u, v)
|
|
w = g.edata['w'][eid].asnumpy()
|
|
assert np.allclose(w, np.zeros(shape=(6, D), dtype=w.dtype))
|
|
|
|
def check_update_routines(readonly):
|
|
g = generate_graph(readonly=readonly)
|
|
g.register_message_func(message_func)
|
|
g.register_reduce_func(reduce_func)
|
|
g.register_apply_node_func(apply_node_func)
|
|
|
|
# send_and_recv
|
|
reduce_msg_shapes.clear()
|
|
u = mx.nd.array([0, 0, 0, 4, 5, 6], dtype='int64')
|
|
v = mx.nd.array([1, 2, 3, 9, 9, 9], dtype='int64')
|
|
g.send_and_recv((u, v))
|
|
assert(reduce_msg_shapes == {(1, 3, D), (3, 1, D)})
|
|
reduce_msg_shapes.clear()
|
|
|
|
# pull
|
|
v = mx.nd.array([1, 2, 3, 9], dtype='int64')
|
|
reduce_msg_shapes.clear()
|
|
g.pull(v)
|
|
assert(reduce_msg_shapes == {(1, 8, D), (3, 1, D)})
|
|
reduce_msg_shapes.clear()
|
|
|
|
# push
|
|
v = mx.nd.array([0, 1, 2, 3], dtype='int64')
|
|
reduce_msg_shapes.clear()
|
|
g.push(v)
|
|
assert(reduce_msg_shapes == {(1, 3, D), (8, 1, D)})
|
|
reduce_msg_shapes.clear()
|
|
|
|
# update_all
|
|
reduce_msg_shapes.clear()
|
|
g.update_all()
|
|
assert(reduce_msg_shapes == {(1, 8, D), (9, 1, D)})
|
|
reduce_msg_shapes.clear()
|
|
|
|
def test_update_routines():
|
|
check_update_routines(True)
|
|
check_update_routines(False)
|
|
|
|
def check_reduce_0deg(readonly):
|
|
if readonly:
|
|
row_idx = []
|
|
col_idx = []
|
|
for i in range(1, 5):
|
|
row_idx.append(i)
|
|
col_idx.append(0)
|
|
ones = np.ones(shape=(len(row_idx)))
|
|
csr = spsp.csr_matrix((ones, (row_idx, col_idx)), shape=(5, 5))
|
|
g = DGLGraph(csr, readonly=True)
|
|
else:
|
|
g = DGLGraph()
|
|
g.add_nodes(5)
|
|
g.add_edge(1, 0)
|
|
g.add_edge(2, 0)
|
|
g.add_edge(3, 0)
|
|
g.add_edge(4, 0)
|
|
def _message(edges):
|
|
return {'m' : edges.src['h']}
|
|
def _reduce(nodes):
|
|
return {'h' : nodes.data['h'] + nodes.mailbox['m'].sum(1)}
|
|
def _init2(shape, dtype, ctx, ids):
|
|
return 2 + mx.nd.zeros(shape, dtype=dtype, ctx=ctx)
|
|
g.set_n_initializer(_init2, 'h')
|
|
old_repr = mx.nd.random.normal(shape=(5, 5))
|
|
g.set_n_repr({'h': old_repr})
|
|
g.update_all(_message, _reduce)
|
|
new_repr = g.ndata['h']
|
|
|
|
assert np.allclose(new_repr[1:].asnumpy(), 2+np.zeros((4, 5)))
|
|
assert np.allclose(new_repr[0].asnumpy(), old_repr.sum(0).asnumpy())
|
|
|
|
def test_reduce_0deg():
|
|
check_reduce_0deg(True)
|
|
check_reduce_0deg(False)
|
|
|
|
def test_recv_0deg_newfld():
|
|
# test recv with 0deg nodes; the reducer also creates a new field
|
|
g = DGLGraph()
|
|
g.add_nodes(2)
|
|
g.add_edge(0, 1)
|
|
def _message(edges):
|
|
return {'m' : edges.src['h']}
|
|
def _reduce(nodes):
|
|
return {'h1' : nodes.data['h'] + mx.nd.sum(nodes.mailbox['m'], 1)}
|
|
def _apply(nodes):
|
|
return {'h1' : nodes.data['h1'] * 2}
|
|
def _init2(shape, dtype, ctx, ids):
|
|
return 2 + mx.nd.zeros(shape=shape, dtype=dtype, ctx=ctx)
|
|
g.register_message_func(_message)
|
|
g.register_reduce_func(_reduce)
|
|
g.register_apply_node_func(_apply)
|
|
# test#1: recv both 0deg and non-0deg nodes
|
|
old = mx.nd.random.normal(shape=(2, 5))
|
|
g.set_n_initializer(_init2, 'h1')
|
|
g.ndata['h'] = old
|
|
g.send((0, 1))
|
|
g.recv([0, 1])
|
|
new = g.ndata.pop('h1')
|
|
# 0deg check: initialized with the func and got applied
|
|
assert np.allclose(new[0].asnumpy(), np.full((5,), 4))
|
|
# non-0deg check
|
|
assert np.allclose(new[1].asnumpy(), mx.nd.sum(old, 0).asnumpy() * 2)
|
|
|
|
# test#2: recv only 0deg node
|
|
old = mx.nd.random.normal(shape=(2, 5))
|
|
g.ndata['h'] = old
|
|
g.ndata['h1'] = mx.nd.full((2, 5), -1) # this is necessary
|
|
g.send((0, 1))
|
|
g.recv(0)
|
|
new = g.ndata.pop('h1')
|
|
# 0deg check: fallback to apply
|
|
assert np.allclose(new[0].asnumpy(), np.full((5,), -2))
|
|
# non-0deg check: not changed
|
|
assert np.allclose(new[1].asnumpy(), np.full((5,), -1))
|
|
|
|
def test_update_all_0deg():
|
|
# test#1
|
|
g = DGLGraph()
|
|
g.add_nodes(5)
|
|
g.add_edge(1, 0)
|
|
g.add_edge(2, 0)
|
|
g.add_edge(3, 0)
|
|
g.add_edge(4, 0)
|
|
def _message(edges):
|
|
return {'m' : edges.src['h']}
|
|
def _reduce(nodes):
|
|
return {'h' : nodes.data['h'] + mx.nd.sum(nodes.mailbox['m'], 1)}
|
|
def _apply(nodes):
|
|
return {'h' : nodes.data['h'] * 2}
|
|
def _init2(shape, dtype, ctx, ids):
|
|
return 2 + mx.nd.zeros(shape, dtype=dtype, ctx=ctx)
|
|
g.set_n_initializer(_init2, 'h')
|
|
old_repr = mx.nd.random.normal(shape=(5, 5))
|
|
g.ndata['h'] = old_repr
|
|
g.update_all(_message, _reduce, _apply)
|
|
new_repr = g.ndata['h']
|
|
# the first row of the new_repr should be the sum of all the node
|
|
# features; while the 0-deg nodes should be initialized by the
|
|
# initializer and applied with UDF.
|
|
assert np.allclose(new_repr[1:].asnumpy(), 2*(2+np.zeros((4,5))))
|
|
assert np.allclose(new_repr[0].asnumpy(), 2 * mx.nd.sum(old_repr, 0).asnumpy())
|
|
|
|
# test#2: graph with no edge
|
|
g = DGLGraph()
|
|
g.add_nodes(5)
|
|
g.set_n_initializer(_init2, 'h')
|
|
g.ndata['h'] = old_repr
|
|
g.update_all(_message, _reduce, _apply)
|
|
new_repr = g.ndata['h']
|
|
# should fallback to apply
|
|
assert np.allclose(new_repr.asnumpy(), 2*old_repr.asnumpy())
|
|
|
|
|
|
def check_pull_0deg(readonly):
|
|
if readonly:
|
|
row_idx = []
|
|
col_idx = []
|
|
row_idx.append(0)
|
|
col_idx.append(1)
|
|
ones = np.ones(shape=(len(row_idx)))
|
|
csr = spsp.csr_matrix((ones, (row_idx, col_idx)), shape=(2, 2))
|
|
g = DGLGraph(csr, readonly=True)
|
|
else:
|
|
g = DGLGraph()
|
|
g.add_nodes(2)
|
|
g.add_edge(0, 1)
|
|
def _message(edges):
|
|
return {'m' : edges.src['h']}
|
|
def _reduce(nodes):
|
|
return {'h' : nodes.mailbox['m'].sum(1)}
|
|
def _apply(nodes):
|
|
return {'h' : nodes.data['h'] * 2}
|
|
def _init2(shape, dtype, ctx, ids):
|
|
return 2 + mx.nd.zeros(shape, dtype=dtype, ctx=ctx)
|
|
g.set_n_initializer(_init2, 'h')
|
|
old_repr = mx.nd.random.normal(shape=(2, 5))
|
|
|
|
# test#1: pull only 0-deg node
|
|
g.ndata['h'] = old_repr
|
|
g.pull(0, _message, _reduce, _apply)
|
|
new_repr = g.ndata['h']
|
|
# 0deg check: equal to apply_nodes
|
|
assert np.allclose(new_repr[0].asnumpy(), old_repr[0].asnumpy() * 2)
|
|
# non-0deg check: untouched
|
|
assert np.allclose(new_repr[1].asnumpy(), old_repr[1].asnumpy())
|
|
|
|
# test#2: pull only non-deg node
|
|
g.ndata['h'] = old_repr
|
|
g.pull(1, _message, _reduce, _apply)
|
|
new_repr = g.ndata['h']
|
|
# 0deg check: untouched
|
|
assert np.allclose(new_repr[0].asnumpy(), old_repr[0].asnumpy())
|
|
# non-0deg check: recved node0 and got applied
|
|
assert np.allclose(new_repr[1].asnumpy(), old_repr[0].asnumpy() * 2)
|
|
|
|
# test#3: pull only both nodes
|
|
g.ndata['h'] = old_repr
|
|
g.pull([0, 1], _message, _reduce, _apply)
|
|
new_repr = g.ndata['h']
|
|
# 0deg check: init and applied
|
|
t = mx.nd.zeros(shape=(2,5)) + 4
|
|
assert np.allclose(new_repr[0].asnumpy(), t.asnumpy())
|
|
# non-0deg check: recv node0 and applied
|
|
assert np.allclose(new_repr[1].asnumpy(), old_repr[0].asnumpy() * 2)
|
|
|
|
def test_pull_0deg():
|
|
check_pull_0deg(True)
|
|
check_pull_0deg(False)
|
|
|
|
if __name__ == '__main__':
|
|
test_batch_setter_getter()
|
|
test_batch_setter_autograd()
|
|
test_batch_send()
|
|
test_batch_recv()
|
|
test_apply_nodes()
|
|
test_apply_edges()
|
|
test_update_routines()
|
|
test_reduce_0deg()
|
|
test_recv_0deg_newfld()
|
|
test_update_all_0deg()
|
|
test_pull_0deg()
|