def test_hetero_conv(agg, idtype):
g = dgl.heterograph(
{
('user', 'follows', 'user'): ([0, 0, 2, 1], [1, 2, 1, 3]),
('user', 'plays', 'game'): ([0, 0, 0, 1, 2], [0, 2, 3, 0, 2]),
('store', 'sells', 'game'): ([0, 0, 1, 1], [0, 3, 1, 2])
},
idtype=idtype,
device=F.ctx())
conv = nn.HeteroGraphConv(
{
'follows': nn.GraphConv(2, 3, allow_zero_in_degree=True),
'plays': nn.GraphConv(2, 4, allow_zero_in_degree=True),
'sells': nn.GraphConv(3, 4, allow_zero_in_degree=True)
}, agg)
conv.initialize(ctx=F.ctx())
print(conv)
uf = F.randn((4, 2))
gf = F.randn((4, 4))
sf = F.randn((2, 3))
h = conv(g, {'user': uf, 'store': sf, 'game': gf})
assert set(h.keys()) == {'user', 'game'}
if agg != 'stack':
assert h['user'].shape == (4, 3)
assert h['game'].shape == (4, 4)
else:
assert h['user'].shape == (4, 1, 3)
assert h['game'].shape == (4, 2, 4)
block = dgl.to_block(g.to(F.cpu()), {
'user': [0, 1, 2, 3],
'game': [0, 1, 2, 3],
'store': []
}).to(F.ctx())
h = conv(block, ({
'user': uf,
'game': gf,
'store': sf
}, {
'user': uf,
'game': gf,
'store': sf[0:0]
}))
assert set(h.keys()) == {'user', 'game'}
if agg != 'stack':
assert h['user'].shape == (4, 3)
assert h['game'].shape == (4, 4)
else:
assert h['user'].shape == (4, 1, 3)
assert h['game'].shape == (4, 2, 4)
h = conv(block, {'user': uf, 'game': gf, 'store': sf})
assert set(h.keys()) == {'user', 'game'}
if agg != 'stack':
assert h['user'].shape == (4, 3)
assert h['game'].shape == (4, 4)
else:
assert h['user'].shape == (4, 1, 3)
assert h['game'].shape == (4, 2, 4)
# test with mod args
class MyMod(mx.gluon.nn.Block):
def __init__(self, s1, s2):
super(MyMod, self).__init__()
self.carg1 = 0
self.s1 = s1
self.s2 = s2
def forward(self, g, h, arg1=None): # mxnet does not support kwargs
if arg1 is not None:
self.carg1 += 1
return F.zeros((g.number_of_dst_nodes(), self.s2))
mod1 = MyMod(2, 3)
mod2 = MyMod(2, 4)
mod3 = MyMod(3, 4)
conv = nn.HeteroGraphConv({
'follows': mod1,
'plays': mod2,
'sells': mod3
}, agg)
conv.initialize(ctx=F.ctx())
mod_args = {'follows': (1, ), 'plays': (1, )}
h = conv(g, {'user': uf, 'store': sf, 'game': gf}, mod_args)
assert mod1.carg1 == 1
assert mod2.carg1 == 1
assert mod3.carg1 == 0
def _test_g(g):
# node subgraph
induced_nodes = [toindex([0, 1, 4]), toindex([0]), toindex([0, 2])]
sub = g.node_subgraph(induced_nodes)
subg = sub.graph
assert subg.number_of_ntypes() == 3
assert subg.number_of_etypes() == 3
assert subg.number_of_nodes(0) == 3
assert subg.number_of_nodes(1) == 1
assert subg.number_of_nodes(2) == 2
assert subg.number_of_edges(R1) == 2
assert subg.number_of_edges(R2) == 1
assert subg.number_of_edges(R3) == 0
adj = subg.adjacency_matrix(R1, True, F.cpu())[0]
assert np.allclose(F.sparse_to_numpy(adj), np.array([[1.], [1.],
[0.]]))
adj = subg.adjacency_matrix(R2, True, F.cpu())[0]
assert np.allclose(F.sparse_to_numpy(adj), np.array([[1., 0.]]))
adj = subg.adjacency_matrix(R3, True, F.cpu())[0]
assert np.allclose(F.sparse_to_numpy(adj), np.array([[0.], [0.]]))
assert len(sub.induced_nodes) == 3
assert _array_equal(sub.induced_nodes[0], induced_nodes[0])
assert _array_equal(sub.induced_nodes[1], induced_nodes[1])
assert _array_equal(sub.induced_nodes[2], induced_nodes[2])
assert len(sub.induced_edges) == 3
assert _array_equal(sub.induced_edges[0], [0, 1])
assert _array_equal(sub.induced_edges[1], [0])
assert _array_equal(sub.induced_edges[2], [])
# node subgraph with empty type graph
induced_nodes = [toindex([0, 1, 4]), toindex([0]), toindex([])]
sub = g.node_subgraph(induced_nodes)
subg = sub.graph
assert subg.number_of_ntypes() == 3
assert subg.number_of_etypes() == 3
assert subg.number_of_nodes(0) == 3
assert subg.number_of_nodes(1) == 1
assert subg.number_of_nodes(2) == 0
assert subg.number_of_edges(R1) == 2
assert subg.number_of_edges(R2) == 0
assert subg.number_of_edges(R3) == 0
adj = subg.adjacency_matrix(R1, True, F.cpu())[0]
assert np.allclose(F.sparse_to_numpy(adj), np.array([[1.], [1.],
[0.]]))
adj = subg.adjacency_matrix(R2, True, F.cpu())[0]
assert np.allclose(F.sparse_to_numpy(adj), np.array([]))
adj = subg.adjacency_matrix(R3, True, F.cpu())[0]
assert np.allclose(F.sparse_to_numpy(adj), np.array([]))
# edge subgraph (preserve_nodes=False)
induced_edges = [toindex([0, 2]), toindex([0]), toindex([0, 1, 2])]
sub = g.edge_subgraph(induced_edges, False)
subg = sub.graph
assert subg.number_of_ntypes() == 3
assert subg.number_of_etypes() == 3
assert subg.number_of_nodes(0) == 2
assert subg.number_of_nodes(1) == 2
assert subg.number_of_nodes(2) == 3
assert subg.number_of_edges(R1) == 2
assert subg.number_of_edges(R2) == 1
assert subg.number_of_edges(R3) == 3
adj = subg.adjacency_matrix(R1, True, F.cpu())[0]
assert np.allclose(F.sparse_to_numpy(adj),
np.array([[1., 0.], [1., 0.]]))
adj = subg.adjacency_matrix(R2, True, F.cpu())[0]
assert np.allclose(F.sparse_to_numpy(adj),
np.array([[1., 0., 0.], [0., 0., 0.]]))
adj = subg.adjacency_matrix(R3, True, F.cpu())[0]
assert np.allclose(F.sparse_to_numpy(adj),
np.array([[0., 1.], [0., 1.], [0., 1.]]))
assert len(sub.induced_nodes) == 3
assert _array_equal(sub.induced_nodes[0], [0, 2])
assert _array_equal(sub.induced_nodes[1], [0, 1])
assert _array_equal(sub.induced_nodes[2], [0, 1, 2])
assert len(sub.induced_edges) == 3
assert _array_equal(sub.induced_edges[0], induced_edges[0])
assert _array_equal(sub.induced_edges[1], induced_edges[1])
assert _array_equal(sub.induced_edges[2], induced_edges[2])
# edge subgraph (preserve_nodes=True)
induced_edges = [toindex([0, 2]), toindex([0]), toindex([0, 1, 2])]
sub = g.edge_subgraph(induced_edges, True)
subg = sub.graph
assert subg.number_of_ntypes() == 3
assert subg.number_of_etypes() == 3
assert subg.number_of_nodes(0) == 5
assert subg.number_of_nodes(1) == 2
assert subg.number_of_nodes(2) == 3
assert subg.number_of_edges(R1) == 2
assert subg.number_of_edges(R2) == 1
assert subg.number_of_edges(R3) == 3
adj = subg.adjacency_matrix(R1, True, F.cpu())[0]
assert np.allclose(
F.sparse_to_numpy(adj),
np.array([[1., 0.], [0., 0.], [1., 0.], [0., 0.], [0., 0.]]))
adj = subg.adjacency_matrix(R2, True, F.cpu())[0]
assert np.allclose(F.sparse_to_numpy(adj),
np.array([[1., 0., 0.], [0., 0., 0.]]))
adj = subg.adjacency_matrix(R3, True, F.cpu())[0]
assert np.allclose(F.sparse_to_numpy(adj),
np.array([[0., 1.], [0., 1.], [0., 1.]]))
assert len(sub.induced_nodes) == 3
assert _array_equal(sub.induced_nodes[0], [0, 1, 2, 3, 4])
assert _array_equal(sub.induced_nodes[1], [0, 1])
assert _array_equal(sub.induced_nodes[2], [0, 1, 2])
assert len(sub.induced_edges) == 3
assert _array_equal(sub.induced_edges[0], induced_edges[0])
assert _array_equal(sub.induced_edges[1], induced_edges[1])
assert _array_equal(sub.induced_edges[2], induced_edges[2])
# edge subgraph with empty induced edges (preserve_nodes=False)
induced_edges = [toindex([0, 2]), toindex([]), toindex([0, 1, 2])]
sub = g.edge_subgraph(induced_edges, False)
subg = sub.graph
assert subg.number_of_ntypes() == 3
assert subg.number_of_etypes() == 3
assert subg.number_of_nodes(0) == 2
assert subg.number_of_nodes(1) == 2
assert subg.number_of_nodes(2) == 3
assert subg.number_of_edges(R1) == 2
assert subg.number_of_edges(R2) == 0
assert subg.number_of_edges(R3) == 3
adj = subg.adjacency_matrix(R1, True, F.cpu())[0]
assert np.allclose(F.sparse_to_numpy(adj),
np.array([[1., 0.], [1., 0.]]))
adj = subg.adjacency_matrix(R2, True, F.cpu())[0]
assert np.allclose(F.sparse_to_numpy(adj),
np.array([[0., 0., 0.], [0., 0., 0.]]))
adj = subg.adjacency_matrix(R3, True, F.cpu())[0]
assert np.allclose(F.sparse_to_numpy(adj),
np.array([[0., 1.], [0., 1.], [0., 1.]]))
assert len(sub.induced_nodes) == 3
assert _array_equal(sub.induced_nodes[0], [0, 2])
assert _array_equal(sub.induced_nodes[1], [0, 1])
assert _array_equal(sub.induced_nodes[2], [0, 1, 2])
assert len(sub.induced_edges) == 3
assert _array_equal(sub.induced_edges[0], induced_edges[0])
assert _array_equal(sub.induced_edges[1], induced_edges[1])
assert _array_equal(sub.induced_edges[2], induced_edges[2])
# edge subgraph with empty induced edges (preserve_nodes=True)
induced_edges = [toindex([0, 2]), toindex([]), toindex([0, 1, 2])]
sub = g.edge_subgraph(induced_edges, True)
subg = sub.graph
assert subg.number_of_ntypes() == 3
assert subg.number_of_etypes() == 3
assert subg.number_of_nodes(0) == 5
assert subg.number_of_nodes(1) == 2
assert subg.number_of_nodes(2) == 3
assert subg.number_of_edges(R1) == 2
assert subg.number_of_edges(R2) == 0
assert subg.number_of_edges(R3) == 3
adj = subg.adjacency_matrix(R1, True, F.cpu())[0]
assert np.allclose(
F.sparse_to_numpy(adj),
np.array([[1., 0.], [0., 0.], [1., 0.], [0., 0.], [0., 0.]]))
adj = subg.adjacency_matrix(R2, True, F.cpu())[0]
assert np.allclose(F.sparse_to_numpy(adj),
np.array([[0., 0., 0.], [0., 0., 0.]]))
adj = subg.adjacency_matrix(R3, True, F.cpu())[0]
assert np.allclose(F.sparse_to_numpy(adj),
np.array([[0., 1.], [0., 1.], [0., 1.]]))
assert len(sub.induced_nodes) == 3
assert _array_equal(sub.induced_nodes[0], [0, 1, 2, 3, 4])
assert _array_equal(sub.induced_nodes[1], [0, 1])
assert _array_equal(sub.induced_nodes[2], [0, 1, 2])
assert len(sub.induced_edges) == 3
assert _array_equal(sub.induced_edges[0], induced_edges[0])
assert _array_equal(sub.induced_edges[1], induced_edges[1])
assert _array_equal(sub.induced_edges[2], induced_edges[2])
def check_dist_graph_hetero(g, num_clients, num_nodes, num_edges):
# Test API
for ntype in num_nodes:
assert ntype in g.ntypes
assert num_nodes[ntype] == g.number_of_nodes(ntype)
for etype in num_edges:
assert etype in g.etypes
assert num_edges[etype] == g.number_of_edges(etype)
etypes = [('n1', 'r1', 'n2'),
('n1', 'r2', 'n3'),
('n2', 'r3', 'n3')]
for i, etype in enumerate(g.canonical_etypes):
assert etype[0] == etypes[i][0]
assert etype[1] == etypes[i][1]
assert etype[2] == etypes[i][2]
assert g.number_of_nodes() == sum([num_nodes[ntype] for ntype in num_nodes])
assert g.number_of_edges() == sum([num_edges[etype] for etype in num_edges])
# Test reading node data
nids = F.arange(0, int(g.number_of_nodes('n1') / 2))
feats1 = g.nodes['n1'].data['feat'][nids]
feats = F.squeeze(feats1, 1)
assert np.all(F.asnumpy(feats == nids))
# Test reading edge data
eids = F.arange(0, int(g.number_of_edges('r1') / 2))
feats1 = g.edges['r1'].data['feat'][eids]
feats = F.squeeze(feats1, 1)
assert np.all(F.asnumpy(feats == eids))
# Test init node data
new_shape = (g.number_of_nodes('n1'), 2)
g.nodes['n1'].data['test1'] = dgl.distributed.DistTensor(new_shape, F.int32)
feats = g.nodes['n1'].data['test1'][nids]
assert np.all(F.asnumpy(feats) == 0)
# create a tensor and destroy a tensor and create it again.
test3 = dgl.distributed.DistTensor(new_shape, F.float32, 'test3', init_func=rand_init)
del test3
test3 = dgl.distributed.DistTensor((g.number_of_nodes('n1'), 3), F.float32, 'test3')
del test3
# add tests for anonymous distributed tensor.
test3 = dgl.distributed.DistTensor(new_shape, F.float32, init_func=rand_init)
data = test3[0:10]
test4 = dgl.distributed.DistTensor(new_shape, F.float32, init_func=rand_init)
del test3
test5 = dgl.distributed.DistTensor(new_shape, F.float32, init_func=rand_init)
assert np.sum(F.asnumpy(test5[0:10] != data)) > 0
# test a persistent tesnor
test4 = dgl.distributed.DistTensor(new_shape, F.float32, 'test4', init_func=rand_init,
persistent=True)
del test4
try:
test4 = dgl.distributed.DistTensor((g.number_of_nodes('n1'), 3), F.float32, 'test4')
raise Exception('')
except:
pass
# Test write data
new_feats = F.ones((len(nids), 2), F.int32, F.cpu())
g.nodes['n1'].data['test1'][nids] = new_feats
feats = g.nodes['n1'].data['test1'][nids]
assert np.all(F.asnumpy(feats) == 1)
# Test metadata operations.
assert len(g.nodes['n1'].data['feat']) == g.number_of_nodes('n1')
assert g.nodes['n1'].data['feat'].shape == (g.number_of_nodes('n1'), 1)
assert g.nodes['n1'].data['feat'].dtype == F.int64
selected_nodes = np.random.randint(0, 100, size=g.number_of_nodes('n1')) > 30
# Test node split
nodes = node_split(selected_nodes, g.get_partition_book(), ntype='n1')
nodes = F.asnumpy(nodes)
# We only have one partition, so the local nodes are basically all nodes in the graph.
local_nids = np.arange(g.number_of_nodes('n1'))
for n in nodes:
assert n in local_nids
print('end')
def emb_init(shape, dtype):
return F.zeros(shape, dtype, F.cpu())
def test_nx_conversion():
# check conversion between networkx and DGLGraph
def _check_nx_feature(nxg, nf, ef):
# check node and edge feature of nxg
# this is used to check to_networkx
num_nodes = len(nxg)
num_edges = nxg.size()
if num_nodes > 0:
node_feat = ddict(list)
for nid, attr in nxg.nodes(data=True):
assert len(attr) == len(nf)
for k in nxg.nodes[nid]:
node_feat[k].append(F.unsqueeze(attr[k], 0))
for k in node_feat:
feat = F.cat(node_feat[k], 0)
assert F.allclose(feat, nf[k])
else:
assert len(nf) == 0
if num_edges > 0:
edge_feat = ddict(lambda: [0] * num_edges)
for u, v, attr in nxg.edges(data=True):
assert len(attr) == len(ef) + 1 # extra id
eid = attr['id']
for k in ef:
edge_feat[k][eid] = F.unsqueeze(attr[k], 0)
for k in edge_feat:
feat = F.cat(edge_feat[k], 0)
assert F.allclose(feat, ef[k])
else:
assert len(ef) == 0
n1 = F.randn((5, 3))
n2 = F.randn((5, 10))
n3 = F.randn((5, 4))
e1 = F.randn((4, 5))
e2 = F.randn((4, 7))
g = dgl.graph([(0, 2), (1, 4), (3, 0), (4, 3)])
g.ndata.update({'n1': n1, 'n2': n2, 'n3': n3})
g.edata.update({'e1': e1, 'e2': e2})
# convert to networkx
nxg = dgl.to_networkx(g, node_attrs=['n1', 'n3'], edge_attrs=['e1', 'e2'])
assert len(nxg) == 5
assert nxg.size() == 4
_check_nx_feature(nxg, {'n1': n1, 'n3': n3}, {'e1': e1, 'e2': e2})
# convert to DGLGraph, nx graph has id in edge feature
# use id feature to test non-tensor copy
g = dgl.graph(nxg, node_attrs=['n1'], edge_attrs=['e1', 'id'])
# check graph size
assert g.number_of_nodes() == 5
assert g.number_of_edges() == 4
# check number of features
# test with existing dglgraph (so existing features should be cleared)
assert len(g.ndata) == 1
assert len(g.edata) == 2
# check feature values
assert F.allclose(g.ndata['n1'], n1)
# with id in nx edge feature, e1 should follow original order
assert F.allclose(g.edata['e1'], e1)
assert F.array_equal(g.edata['id'], F.copy_to(F.arange(0, 4), F.cpu()))
# test conversion after modifying DGLGraph
# TODO(minjie): enable after mutation is supported
#g.pop_e_repr('id') # pop id so we don't need to provide id when adding edges
#new_n = F.randn((2, 3))
#new_e = F.randn((3, 5))
#g.add_nodes(2, data={'n1': new_n})
## add three edges, one is a multi-edge
#g.add_edges([3, 6, 0], [4, 5, 2], data={'e1': new_e})
#n1 = F.cat((n1, new_n), 0)
#e1 = F.cat((e1, new_e), 0)
## convert to networkx again
#nxg = g.to_networkx(node_attrs=['n1'], edge_attrs=['e1'])
#assert len(nxg) == 7
#assert nxg.size() == 7
#_check_nx_feature(nxg, {'n1': n1}, {'e1': e1})
# now test convert from networkx without id in edge feature
# first pop id in edge feature
for _, _, attr in nxg.edges(data=True):
attr.pop('id')
# test with a new graph
g = dgl.graph(nxg, node_attrs=['n1'], edge_attrs=['e1'])
# check graph size
assert g.number_of_nodes() == 5
assert g.number_of_edges() == 4
# check number of features
assert len(g.ndata) == 1
assert len(g.edata) == 1
# check feature values
assert F.allclose(g.ndata['n1'], n1)
# edge feature order follows nxg.edges()
edge_feat = []
for _, _, attr in nxg.edges(data=True):
edge_feat.append(F.unsqueeze(attr['e1'], 0))
edge_feat = F.cat(edge_feat, 0)
assert F.allclose(g.edata['e1'], edge_feat)
# Test converting from a networkx graph whose nodes are
# not labeled with consecutive-integers.
nxg = nx.cycle_graph(5)
nxg.remove_nodes_from([0, 4])
for u in nxg.nodes():
nxg.node[u]['h'] = F.tensor([u])
for u, v, d in nxg.edges(data=True):
d['h'] = F.tensor([u, v])
g = dgl.DGLGraph()
g.from_networkx(nxg, node_attrs=['h'], edge_attrs=['h'])
assert g.number_of_nodes() == 3
assert g.number_of_edges() == 4
assert g.has_edge_between(0, 1)
assert g.has_edge_between(1, 2)
assert F.allclose(g.ndata['h'], F.tensor([[1.], [2.], [3.]]))
assert F.allclose(g.edata['h'],
F.tensor([[1., 2.], [1., 2.], [2., 3.], [2., 3.]]))
def check_dist_graph(g, num_clients, num_nodes, num_edges):
# Test API
assert g.number_of_nodes() == num_nodes
assert g.number_of_edges() == num_edges
# Test reading node data
nids = F.arange(0, int(g.number_of_nodes() / 2))
feats1 = g.ndata['features'][nids]
feats = F.squeeze(feats1, 1)
assert np.all(F.asnumpy(feats == nids))
# Test reading edge data
eids = F.arange(0, int(g.number_of_edges() / 2))
feats1 = g.edata['features'][eids]
feats = F.squeeze(feats1, 1)
assert np.all(F.asnumpy(feats == eids))
# Test init node data
new_shape = (g.number_of_nodes(), 2)
g.ndata['test1'] = dgl.distributed.DistTensor(new_shape, F.int32)
feats = g.ndata['test1'][nids]
assert np.all(F.asnumpy(feats) == 0)
# reference to a one that exists
test2 = dgl.distributed.DistTensor(new_shape, F.float32, 'test2', init_func=rand_init)
test3 = dgl.distributed.DistTensor(new_shape, F.float32, 'test2')
assert np.all(F.asnumpy(test2[nids]) == F.asnumpy(test3[nids]))
# create a tensor and destroy a tensor and create it again.
test3 = dgl.distributed.DistTensor(new_shape, F.float32, 'test3', init_func=rand_init)
del test3
test3 = dgl.distributed.DistTensor((g.number_of_nodes(), 3), F.float32, 'test3')
del test3
# add tests for anonymous distributed tensor.
test3 = dgl.distributed.DistTensor(new_shape, F.float32, init_func=rand_init)
data = test3[0:10]
test4 = dgl.distributed.DistTensor(new_shape, F.float32, init_func=rand_init)
del test3
test5 = dgl.distributed.DistTensor(new_shape, F.float32, init_func=rand_init)
assert np.sum(F.asnumpy(test5[0:10] != data)) > 0
# test a persistent tesnor
test4 = dgl.distributed.DistTensor(new_shape, F.float32, 'test4', init_func=rand_init,
persistent=True)
del test4
try:
test4 = dgl.distributed.DistTensor((g.number_of_nodes(), 3), F.float32, 'test4')
raise Exception('')
except:
pass
# Test write data
new_feats = F.ones((len(nids), 2), F.int32, F.cpu())
g.ndata['test1'][nids] = new_feats
feats = g.ndata['test1'][nids]
assert np.all(F.asnumpy(feats) == 1)
# Test metadata operations.
assert len(g.ndata['features']) == g.number_of_nodes()
assert g.ndata['features'].shape == (g.number_of_nodes(), 1)
assert g.ndata['features'].dtype == F.int64
assert g.node_attr_schemes()['features'].dtype == F.int64
assert g.node_attr_schemes()['test1'].dtype == F.int32
assert g.node_attr_schemes()['features'].shape == (1,)
selected_nodes = np.random.randint(0, 100, size=g.number_of_nodes()) > 30
# Test node split
nodes = node_split(selected_nodes, g.get_partition_book())
nodes = F.asnumpy(nodes)
# We only have one partition, so the local nodes are basically all nodes in the graph.
local_nids = np.arange(g.number_of_nodes())
for n in nodes:
assert n in local_nids
print('end')
def test_knn_cpu(algorithm, dist):
x = th.randn(8, 3).to(F.cpu())
kg = dgl.nn.KNNGraph(3)
if dist == 'euclidean':
d = th.cdist(x, x).to(F.cpu())
else:
x = x + th.randn(1).item()
tmp_x = x / (1e-5 + F.sqrt(F.sum(x * x, dim=1, keepdims=True)))
d = 1 - F.matmul(tmp_x, tmp_x.T).to(F.cpu())
def check_knn(g, x, start, end, k):
assert g.device == x.device
for v in range(start, end):
src, _ = g.in_edges(v)
src = set(src.numpy())
i = v - start
src_ans = set(th.topk(d[start:end, start:end][i], k, largest=False)[1].numpy() + start)
assert src == src_ans
# check knn with 2d input
g = kg(x, algorithm, dist)
check_knn(g, x, 0, 8, 3)
# check knn with 3d input
g = kg(x.view(2, 4, 3), algorithm, dist)
check_knn(g, x, 0, 4, 3)
check_knn(g, x, 4, 8, 3)
# check segmented knn
kg = dgl.nn.SegmentedKNNGraph(3)
g = kg(x, [3, 5], algorithm, dist)
check_knn(g, x, 0, 3, 3)
check_knn(g, x, 3, 8, 3)
# check k > num_points
kg = dgl.nn.KNNGraph(10)
with pytest.warns(DGLWarning):
g = kg(x, algorithm, dist)
check_knn(g, x, 0, 8, 8)
with pytest.warns(DGLWarning):
g = kg(x.view(2, 4, 3), algorithm, dist)
check_knn(g, x, 0, 4, 4)
check_knn(g, x, 4, 8, 4)
kg = dgl.nn.SegmentedKNNGraph(5)
with pytest.warns(DGLWarning):
g = kg(x, [3, 5], algorithm, dist)
check_knn(g, x, 0, 3, 3)
check_knn(g, x, 3, 8, 3)
# check k == 0
kg = dgl.nn.KNNGraph(0)
with pytest.raises(DGLError):
g = kg(x, algorithm, dist)
kg = dgl.nn.SegmentedKNNGraph(0)
with pytest.raises(DGLError):
g = kg(x, [3, 5], algorithm, dist)
# check empty
x_empty = th.tensor([])
kg = dgl.nn.KNNGraph(3)
with pytest.raises(DGLError):
g = kg(x_empty, algorithm, dist)
kg = dgl.nn.SegmentedKNNGraph(3)
with pytest.raises(DGLError):
g = kg(x_empty, [3, 5], algorithm, dist)
import os
import time
import dgl
import backend as F
import unittest, pytest
from numpy.testing import assert_array_equal
INTEGER = 2
STR = 'hello world!'
HELLO_SERVICE_ID = 901231
TENSOR = F.zeros((10, 10), F.int64, F.cpu())
def foo(x, y):
assert x == 123
assert y == "abc"
class MyRequest(dgl.distributed.Request):
def __init__(self):
self.x = 123
self.y = "abc"
self.z = F.randn((3, 4))
self.foo = foo
def __getstate__(self):
return self.x, self.y, self.z, self.foo
def __setstate__(self, state):
def test_hetero_conv(agg, idtype):
g = dgl.heterograph(
{
('user', 'follows', 'user'): ([0, 0, 2, 1], [1, 2, 1, 3]),
('user', 'plays', 'game'): ([0, 0, 0, 1, 2], [0, 2, 3, 0, 2]),
('store', 'sells', 'game'): ([0, 0, 1, 1], [0, 3, 1, 2])
},
idtype=idtype,
device=F.ctx())
conv = nn.HeteroGraphConv(
{
'follows': nn.GraphConv(2, 3, allow_zero_in_degree=True),
'plays': nn.GraphConv(2, 4, allow_zero_in_degree=True),
'sells': nn.GraphConv(3, 4, allow_zero_in_degree=True)
}, agg)
uf = F.randn((4, 2))
gf = F.randn((4, 4))
sf = F.randn((2, 3))
h = conv(g, {'user': uf, 'store': sf, 'game': gf})
assert set(h.keys()) == {'user', 'game'}
if agg != 'stack':
assert h['user'].shape == (4, 3)
assert h['game'].shape == (4, 4)
else:
assert h['user'].shape == (4, 1, 3)
assert h['game'].shape == (4, 2, 4)
block = dgl.to_block(g.to(F.cpu()), {
'user': [0, 1, 2, 3],
'game': [0, 1, 2, 3],
'store': []
}).to(F.ctx())
h = conv(block, ({
'user': uf,
'game': gf,
'store': sf
}, {
'user': uf,
'game': gf,
'store': sf[0:0]
}))
assert set(h.keys()) == {'user', 'game'}
if agg != 'stack':
assert h['user'].shape == (4, 3)
assert h['game'].shape == (4, 4)
else:
assert h['user'].shape == (4, 1, 3)
assert h['game'].shape == (4, 2, 4)
h = conv(block, {'user': uf, 'game': gf, 'store': sf})
assert set(h.keys()) == {'user', 'game'}
if agg != 'stack':
assert h['user'].shape == (4, 3)
assert h['game'].shape == (4, 4)
else:
assert h['user'].shape == (4, 1, 3)
assert h['game'].shape == (4, 2, 4)
# test with mod args
class MyMod(tf.keras.layers.Layer):
def __init__(self, s1, s2):
super(MyMod, self).__init__()
self.carg1 = 0
self.carg2 = 0
self.s1 = s1
self.s2 = s2
def call(self, g, h, arg1=None, *, arg2=None):
if arg1 is not None:
self.carg1 += 1
if arg2 is not None:
self.carg2 += 1
return tf.zeros((g.number_of_dst_nodes(), self.s2))
mod1 = MyMod(2, 3)
mod2 = MyMod(2, 4)
mod3 = MyMod(3, 4)
conv = nn.HeteroGraphConv({
'follows': mod1,
'plays': mod2,
'sells': mod3
}, agg)
mod_args = {'follows': (1, ), 'plays': (1, )}
mod_kwargs = {'sells': {'arg2': 'abc'}}
h = conv(g, {
'user': uf,
'game': gf,
'store': sf
},
mod_args=mod_args,
mod_kwargs=mod_kwargs)
assert mod1.carg1 == 1
assert mod1.carg2 == 0
assert mod2.carg1 == 1
assert mod2.carg2 == 0
assert mod3.carg1 == 0
assert mod3.carg2 == 1
def init_zero_func(shape, dtype):
return F.zeros(shape, dtype, F.cpu())
def check_dist_graph(g, num_clients, num_nodes, num_edges):
# Test API
assert g.number_of_nodes() == num_nodes
assert g.number_of_edges() == num_edges
# Test reading node data
nids = F.arange(0, int(g.number_of_nodes() / 2))
feats1 = g.ndata['features'][nids]
feats = F.squeeze(feats1, 1)
assert np.all(F.asnumpy(feats == nids))
# Test reading edge data
eids = F.arange(0, int(g.number_of_edges() / 2))
feats1 = g.edata['features'][eids]
feats = F.squeeze(feats1, 1)
assert np.all(F.asnumpy(feats == eids))
# Test init node data
new_shape = (g.number_of_nodes(), 2)
g.ndata['test1'] = dgl.distributed.DistTensor(new_shape, F.int32)
feats = g.ndata['test1'][nids]
assert np.all(F.asnumpy(feats) == 0)
# reference to a one that exists
test2 = dgl.distributed.DistTensor(new_shape, F.float32, 'test2', init_func=rand_init)
test3 = dgl.distributed.DistTensor(new_shape, F.float32, 'test2')
assert np.all(F.asnumpy(test2[nids]) == F.asnumpy(test3[nids]))
# create a tensor and destroy a tensor and create it again.
test3 = dgl.distributed.DistTensor(new_shape, F.float32, 'test3', init_func=rand_init)
del test3
test3 = dgl.distributed.DistTensor((g.number_of_nodes(), 3), F.float32, 'test3')
del test3
# test a persistent tesnor
test4 = dgl.distributed.DistTensor(new_shape, F.float32, 'test4', init_func=rand_init,
persistent=True)
del test4
try:
test4 = dgl.distributed.DistTensor((g.number_of_nodes(), 3), F.float32, 'test4')
raise Exception('')
except:
pass
# Test sparse emb
try:
emb = DistEmbedding(g.number_of_nodes(), 1, 'emb1', emb_init)
lr = 0.001
optimizer = SparseAdagrad([emb], lr=lr)
with F.record_grad():
feats = emb(nids)
assert np.all(F.asnumpy(feats) == np.zeros((len(nids), 1)))
loss = F.sum(feats + 1, 0)
loss.backward()
optimizer.step()
feats = emb(nids)
if num_clients == 1:
assert_almost_equal(F.asnumpy(feats), np.ones((len(nids), 1)) * -lr)
rest = np.setdiff1d(np.arange(g.number_of_nodes()), F.asnumpy(nids))
feats1 = emb(rest)
assert np.all(F.asnumpy(feats1) == np.zeros((len(rest), 1)))
policy = dgl.distributed.PartitionPolicy('node', g.get_partition_book())
grad_sum = dgl.distributed.DistTensor((g.number_of_nodes(),), F.float32,
'emb1_sum', policy)
if num_clients == 1:
assert np.all(F.asnumpy(grad_sum[nids]) == np.ones((len(nids), 1)) * num_clients)
assert np.all(F.asnumpy(grad_sum[rest]) == np.zeros((len(rest), 1)))
emb = DistEmbedding(g.number_of_nodes(), 1, 'emb2', emb_init)
with F.no_grad():
feats1 = emb(nids)
assert np.all(F.asnumpy(feats1) == 0)
optimizer = SparseAdagrad([emb], lr=lr)
with F.record_grad():
feats1 = emb(nids)
feats2 = emb(nids)
feats = F.cat([feats1, feats2], 0)
assert np.all(F.asnumpy(feats) == np.zeros((len(nids) * 2, 1)))
loss = F.sum(feats + 1, 0)
loss.backward()
optimizer.step()
with F.no_grad():
feats = emb(nids)
if num_clients == 1:
assert_almost_equal(F.asnumpy(feats), np.ones((len(nids), 1)) * math.sqrt(2) * -lr)
rest = np.setdiff1d(np.arange(g.number_of_nodes()), F.asnumpy(nids))
feats1 = emb(rest)
assert np.all(F.asnumpy(feats1) == np.zeros((len(rest), 1)))
except NotImplementedError as e:
pass
# Test write data
new_feats = F.ones((len(nids), 2), F.int32, F.cpu())
g.ndata['test1'][nids] = new_feats
feats = g.ndata['test1'][nids]
assert np.all(F.asnumpy(feats) == 1)
# Test metadata operations.
assert len(g.ndata['features']) == g.number_of_nodes()
assert g.ndata['features'].shape == (g.number_of_nodes(), 1)
assert g.ndata['features'].dtype == F.int64
assert g.node_attr_schemes()['features'].dtype == F.int64
assert g.node_attr_schemes()['test1'].dtype == F.int32
assert g.node_attr_schemes()['features'].shape == (1,)
selected_nodes = np.random.randint(0, 100, size=g.number_of_nodes()) > 30
# Test node split
nodes = node_split(selected_nodes, g.get_partition_book())
nodes = F.asnumpy(nodes)
# We only have one partition, so the local nodes are basically all nodes in the graph.
local_nids = np.arange(g.number_of_nodes())
for n in nodes:
assert n in local_nids
print('end')
def test_random_walk():
g1 = dgl.heterograph({
('user', 'follow', 'user'): [(0, 1), (1, 2), (2, 0)]
})
g2 = dgl.heterograph({
('user', 'follow', 'user'): [(0, 1), (1, 2), (1, 3), (2, 0), (3, 0)]
})
g3 = dgl.heterograph({
('user', 'follow', 'user'): [(0, 1), (1, 2), (2, 0)],
('user', 'view', 'item'): [(0, 0), (1, 1), (2, 2)],
('item', 'viewed-by', 'user'): [(0, 0), (1, 1), (2, 2)]
})
g4 = dgl.heterograph({
('user', 'follow', 'user'): [(0, 1), (1, 2), (1, 3), (2, 0), (3, 0)],
('user', 'view', 'item'): [(0, 0), (0, 1), (1, 1), (2, 2), (3, 2),
(3, 1)],
('item', 'viewed-by', 'user'): [(0, 0), (1, 0), (1, 1), (2, 2), (2, 3),
(1, 3)]
})
g2.edata['p'] = F.tensor([3, 0, 3, 3, 3], dtype=F.float32)
g2.edata['p2'] = F.tensor([[3], [0], [3], [3], [3]], dtype=F.float32)
g4.edges['follow'].data['p'] = F.tensor([3, 0, 3, 3, 3], dtype=F.float32)
g4.edges['viewed-by'].data['p'] = F.tensor([1, 1, 1, 1, 1, 1],
dtype=F.float32)
traces, ntypes = dgl.sampling.random_walk(g1, [0, 1, 2, 0, 1, 2], length=4)
check_random_walk(g1, ['follow'] * 4, traces, ntypes)
traces, ntypes = dgl.sampling.random_walk(g1, [0, 1, 2, 0, 1, 2],
length=4,
restart_prob=0.)
check_random_walk(g1, ['follow'] * 4, traces, ntypes)
traces, ntypes = dgl.sampling.random_walk(g1, [0, 1, 2, 0, 1, 2],
length=4,
restart_prob=F.zeros((4, ),
F.float32,
F.cpu()))
check_random_walk(g1, ['follow'] * 4, traces, ntypes)
traces, ntypes = dgl.sampling.random_walk(g1, [0, 1, 2, 0, 1, 2],
length=5,
restart_prob=F.tensor(
[0, 0, 0, 0, 1],
dtype=F.float32))
check_random_walk(g1, ['follow'] * 4, F.slice_axis(traces, 1, 0, 5),
F.slice_axis(ntypes, 0, 0, 5))
assert (F.asnumpy(traces)[:, 5] == -1).all()
traces, ntypes = dgl.sampling.random_walk(g2, [0, 1, 2, 3, 0, 1, 2, 3],
length=4)
check_random_walk(g2, ['follow'] * 4, traces, ntypes)
traces, ntypes = dgl.sampling.random_walk(g2, [0, 1, 2, 3, 0, 1, 2, 3],
length=4,
prob='p')
check_random_walk(g2, ['follow'] * 4, traces, ntypes, 'p')
try:
traces, ntypes = dgl.sampling.random_walk(g2, [0, 1, 2, 3, 0, 1, 2, 3],
length=4,
prob='p2')
fail = False
except dgl.DGLError:
fail = True
assert fail
metapath = ['follow', 'view', 'viewed-by'] * 2
traces, ntypes = dgl.sampling.random_walk(g3, [0, 1, 2, 0, 1, 2],
metapath=metapath)
check_random_walk(g3, metapath, traces, ntypes)
metapath = ['follow', 'view', 'viewed-by'] * 2
traces, ntypes = dgl.sampling.random_walk(g4, [0, 1, 2, 3, 0, 1, 2, 3],
metapath=metapath)
check_random_walk(g4, metapath, traces, ntypes)
metapath = ['follow', 'view', 'viewed-by'] * 2
traces, ntypes = dgl.sampling.random_walk(g4, [0, 1, 2, 3, 0, 1, 2, 3],
metapath=metapath,
prob='p')
check_random_walk(g4, metapath, traces, ntypes, 'p')
traces, ntypes = dgl.sampling.random_walk(g4, [0, 1, 2, 3, 0, 1, 2, 3],
metapath=metapath,
prob='p',
restart_prob=0.)
check_random_walk(g4, metapath, traces, ntypes, 'p')
traces, ntypes = dgl.sampling.random_walk(g4, [0, 1, 2, 3, 0, 1, 2, 3],
metapath=metapath,
prob='p',
restart_prob=F.zeros((6, ),
F.float32,
F.cpu()))
check_random_walk(g4, metapath, traces, ntypes, 'p')
traces, ntypes = dgl.sampling.random_walk(
g4, [0, 1, 2, 3, 0, 1, 2, 3],
metapath=metapath + ['follow'],
prob='p',
restart_prob=F.tensor([0, 0, 0, 0, 0, 0, 1], F.float32))
check_random_walk(g4, metapath, traces[:, :7], ntypes[:7], 'p')
assert (F.asnumpy(traces[:, 7]) == -1).all()
import scipy as sp
import dgl
from dgl import utils
from dgl.contrib import KVServer
from dgl.contrib import KVClient
from numpy.testing import assert_array_equal
import os
import time
num_entries = 10
dim_size = 3
server_namebook = {0:[0, '127.0.0.1', 30070, 1]}
data_0 = F.zeros((num_entries, dim_size), F.float32, F.cpu())
g2l_0 = F.arange(0, num_entries)
partition_0 = F.zeros(num_entries, F.int64, F.cpu())
data_1 = F.zeros((num_entries*2, dim_size), F.float32, F.cpu())
g2l_1 = F.arange(0, num_entries*2)
partition_1 = F.zeros(num_entries*2, F.int64, F.cpu())
data_3 = F.zeros((num_entries, dim_size), F.int64, F.cpu())
data_4 = F.zeros((num_entries, dim_size), F.float64, F.cpu())
data_5 = F.zeros((num_entries, dim_size), F.int32, F.cpu())
def start_server():
my_server = KVServer(server_id=0, server_namebook=server_namebook, num_client=1)
my_server.set_global2local(name='data_0', global2local=g2l_0)
my_server.set_global2local(name='data_1', global2local=g2l_1)
