def test_glob_att_pool():
g = dgl.DGLGraph(nx.path_graph(10))
ctx = F.ctx()
gap = nn.GlobalAttentionPooling(gluon.nn.Dense(1), gluon.nn.Dense(10))
gap.initialize(ctx=ctx)
print(gap)
# test#1: basic
h0 = F.randn((g.number_of_nodes(), 5))
h1 = gap(g, h0)
assert h1.shape[0] == 10 and h1.ndim == 1
# test#2: batched graph
bg = dgl.batch([g, g, g, g])
h0 = F.randn((bg.number_of_nodes(), 5))
h1 = gap(bg, h0)
assert h1.shape[0] == 4 and h1.shape[1] == 10 and h1.ndim == 2
def test_dense_graph_conv():
ctx = F.ctx()
g = dgl.DGLGraph(sp.sparse.random(100, 100, density=0.1), readonly=True)
adj = g.adjacency_matrix(ctx=ctx).to_dense()
conv = nn.GraphConv(5, 2, norm=False, bias=True)
dense_conv = nn.DenseGraphConv(5, 2, norm=False, bias=True)
dense_conv.weight.data = conv.weight.data
dense_conv.bias.data = conv.bias.data
feat = F.randn((100, 5))
if F.gpu_ctx():
conv = conv.to(ctx)
dense_conv = dense_conv.to(ctx)
feat = feat.to(ctx)
out_conv = conv(g, feat)
out_dense_conv = dense_conv(adj, feat)
assert F.allclose(out_conv, out_dense_conv)
def test_graph_conv():
g = dgl.DGLGraph(nx.path_graph(3))
ctx = F.ctx()
adj = tf.sparse.to_dense(tf.sparse.reorder(g.adjacency_matrix(ctx=ctx)))
conv = nn.GraphConv(5, 2, norm='none', bias=True)
# conv = conv
print(conv)
# test#1: basic
h0 = F.ones((3, 5))
h1 = conv(g, h0)
assert len(g.ndata) == 0
assert len(g.edata) == 0
assert F.allclose(h1, _AXWb(adj, h0, conv.weight, conv.bias))
# test#2: more-dim
h0 = F.ones((3, 5, 5))
h1 = conv(g, h0)
assert len(g.ndata) == 0
assert len(g.edata) == 0
assert F.allclose(h1, _AXWb(adj, h0, conv.weight, conv.bias))
conv = nn.GraphConv(5, 2)
# conv = conv
# test#3: basic
h0 = F.ones((3, 5))
h1 = conv(g, h0)
assert len(g.ndata) == 0
assert len(g.edata) == 0
# test#4: basic
h0 = F.ones((3, 5, 5))
h1 = conv(g, h0)
assert len(g.ndata) == 0
assert len(g.edata) == 0
conv = nn.GraphConv(5, 2)
# conv = conv
# test#3: basic
h0 = F.ones((3, 5))
h1 = conv(g, h0)
assert len(g.ndata) == 0
assert len(g.edata) == 0
# test#4: basic
h0 = F.ones((3, 5, 5))
h1 = conv(g, h0)
assert len(g.ndata) == 0
assert len(g.edata) == 0
def test_dense_graph_conv():
ctx = F.ctx()
g = dgl.DGLGraph(sp.sparse.random(100, 100, density=0.3), readonly=True)
adj = g.adjacency_matrix(ctx=ctx).tostype('default')
conv = nn.GraphConv(5, 2, norm=False, bias=True)
dense_conv = nn.DenseGraphConv(5, 2, norm=False, bias=True)
conv.initialize(ctx=ctx)
dense_conv.initialize(ctx=ctx)
dense_conv.weight.set_data(
conv.weight.data())
dense_conv.bias.set_data(
conv.bias.data())
feat = F.randn((100, 5))
out_conv = conv(g, feat)
out_dense_conv = dense_conv(adj, feat)
assert F.allclose(out_conv, out_dense_conv)
def test_atomic_conv():
g = dgl.DGLGraph(sp.sparse.random(100, 100, density=0.1), readonly=True)
aconv = nn.AtomicConv(interaction_cutoffs=F.tensor([12.0, 12.0]),
rbf_kernel_means=F.tensor([0.0, 2.0]),
rbf_kernel_scaling=F.tensor([4.0, 4.0]),
features_to_use=F.tensor([6.0, 8.0]))
ctx = F.ctx()
if F.gpu_ctx():
aconv = aconv.to(ctx)
feat = F.randn((100, 1))
dist = F.randn((g.number_of_edges(), 1))
h = aconv(g, feat, dist)
# current we only do shape check
assert h.shape[-1] == 4
def test_edge_softmax():
# Basic
g = dgl.DGLGraph(nx.path_graph(3)).to(F.ctx())
edata = F.ones((g.number_of_edges(), 1))
a = nn.edge_softmax(g, edata)
assert len(g.ndata) == 0
assert len(g.edata) == 0
assert np.allclose(a.asnumpy(), uniform_attention(g, a.shape).asnumpy(),
1e-4, 1e-4)
# Test higher dimension case
edata = F.ones((g.number_of_edges(), 3, 1))
a = nn.edge_softmax(g, edata)
assert len(g.ndata) == 0
assert len(g.edata) == 0
assert np.allclose(a.asnumpy(), uniform_attention(g, a.shape).asnumpy(),
1e-4, 1e-4)
def test_edge_predictor(op):
ctx = F.ctx()
num_pairs = 3
in_feats = 4
out_feats = 5
h_src = th.randn((num_pairs, in_feats)).to(ctx)
h_dst = th.randn((num_pairs, in_feats)).to(ctx)
pred = nn.EdgePredictor(op)
if op in ['dot', 'cos']:
assert pred(h_src, h_dst).shape == (num_pairs, 1)
elif op == 'ele':
assert pred(h_src, h_dst).shape == (num_pairs, in_feats)
else:
assert pred(h_src, h_dst).shape == (num_pairs, 2 * in_feats)
pred = nn.EdgePredictor(op, in_feats, out_feats, bias=True).to(ctx)
assert pred(h_src, h_dst).shape == (num_pairs, out_feats)
def test_hetero_embedding(out_dim):
layer = nn.HeteroEmbedding({
'user': 2,
('user', 'follows', 'user'): 3
}, out_dim)
layer = layer.to(F.ctx())
embeds = layer.weight
assert embeds['user'].shape == (2, out_dim)
assert embeds[('user', 'follows', 'user')].shape == (3, out_dim)
embeds = layer({
'user': F.tensor([0], dtype=F.int64),
('user', 'follows', 'user'): F.tensor([0, 2], dtype=F.int64)
})
assert embeds['user'].shape == (1, out_dim)
assert embeds[('user', 'follows', 'user')].shape == (2, out_dim)
def test_dense_sage_conv():
ctx = F.ctx()
g = dgl.DGLGraph(sp.sparse.random(100, 100, density=0.1), readonly=True)
adj = g.adjacency_matrix(ctx=ctx).tostype('default')
sage = nn.SAGEConv(5, 2, 'gcn')
dense_sage = nn.DenseSAGEConv(5, 2)
sage.initialize(ctx=ctx)
dense_sage.initialize(ctx=ctx)
dense_sage.fc.weight.set_data(
sage.fc_neigh.weight.data())
dense_sage.fc.bias.set_data(
sage.fc_neigh.bias.data())
feat = F.randn((100, 5))
out_sage = sage(g, feat)
out_dense_sage = dense_sage(adj, feat)
assert F.allclose(out_sage, out_dense_sage)
def test_dense_sage_conv():
ctx = F.ctx()
g = dgl.DGLGraph(sp.sparse.random(100, 100, density=0.1), readonly=True)
adj = g.adjacency_matrix(ctx=ctx).to_dense()
sage = nn.SAGEConv(5, 2, 'gcn',)
dense_sage = nn.DenseSAGEConv(5, 2)
dense_sage.fc.weight.data = sage.fc_neigh.weight.data
dense_sage.fc.bias.data = sage.fc_neigh.bias.data
feat = F.randn((100, 5))
if F.gpu_ctx():
sage = sage.to(ctx)
dense_sage = dense_sage.to(ctx)
feat = feat.to(ctx)
out_sage = sage(g, feat)
out_dense_sage = dense_sage(adj, feat)
assert F.allclose(out_sage, out_dense_sage)
def test_rgcn():
ctx = F.ctx()
etype = []
g = dgl.DGLGraph(sp.sparse.random(100, 100, density=0.1), readonly=True)
# 5 etypes
R = 5
for i in range(g.number_of_edges()):
etype.append(i % 5)
B = 2
I = 10
O = 8
rgc_basis = nn.RelGraphConv(I, O, R, "basis", B).to(ctx)
h = th.randn((100, I)).to(ctx)
r = th.tensor(etype).to(ctx)
h_new = rgc_basis(g, h, r)
assert list(h_new.shape) == [100, O]
rgc_bdd = nn.RelGraphConv(I, O, R, "bdd", B).to(ctx)
h = th.randn((100, I)).to(ctx)
r = th.tensor(etype).to(ctx)
h_new = rgc_bdd(g, h, r)
assert list(h_new.shape) == [100, O]
# with norm
norm = th.zeros((g.number_of_edges(), 1)).to(ctx)
rgc_basis = nn.RelGraphConv(I, O, R, "basis", B).to(ctx)
h = th.randn((100, I)).to(ctx)
r = th.tensor(etype).to(ctx)
h_new = rgc_basis(g, h, r, norm)
assert list(h_new.shape) == [100, O]
rgc_bdd = nn.RelGraphConv(I, O, R, "bdd", B).to(ctx)
h = th.randn((100, I)).to(ctx)
r = th.tensor(etype).to(ctx)
h_new = rgc_bdd(g, h, r, norm)
assert list(h_new.shape) == [100, O]
# id input
rgc_basis = nn.RelGraphConv(I, O, R, "basis", B).to(ctx)
h = th.randint(0, I, (100, )).to(ctx)
r = th.tensor(etype).to(ctx)
h_new = rgc_basis(g, h, r)
assert list(h_new.shape) == [100, O]
def test_set2set():
g = dgl.DGLGraph(nx.path_graph(10))
ctx = F.ctx()
s2s = nn.Set2Set(5, 3, 3) # hidden size 5, 3 iters, 3 layers
s2s.initialize(ctx=ctx)
print(s2s)
# test#1: basic
h0 = F.randn((g.number_of_nodes(), 5))
h1 = s2s(g, h0)
assert h1.shape[0] == 1 and h1.shape[1] == 10 and h1.ndim == 2
# test#2: batched graph
bg = dgl.batch([g, g, g])
h0 = F.randn((bg.number_of_nodes(), 5))
h1 = s2s(bg, h0)
assert h1.shape[0] == 3 and h1.shape[1] == 10 and h1.ndim == 2
def test_sgc_conv():
ctx = F.ctx()
g = dgl.DGLGraph(sp.sparse.random(100, 100, density=0.1), readonly=True)
# not cached
sgc = nn.SGConv(5, 10, 3)
feat = F.randn((100, 5))
sgc = sgc.to(ctx)
h = sgc(g, feat)
assert h.shape[-1] == 10
# cached
sgc = nn.SGConv(5, 10, 3, True)
sgc = sgc.to(ctx)
h_0 = sgc(g, feat)
h_1 = sgc(g, feat + 1)
assert F.allclose(h_0, h_1)
assert h_0.shape[-1] == 10
def test_ke_score_funcs():
ctx = F.ctx()
num_edges = 30
num_rels = 3
nfeats = 4
h_src = th.randn((num_edges, nfeats)).to(ctx)
h_dst = th.randn((num_edges, nfeats)).to(ctx)
rels = th.randint(low=0, high=num_rels, size=(num_edges, )).to(ctx)
score_func = nn.TransE(num_rels=num_rels, feats=nfeats).to(ctx)
score_func.reset_parameters()
score_func(h_src, h_dst, rels).shape == (num_edges)
score_func = nn.TransR(num_rels=num_rels, rfeats=nfeats - 1,
nfeats=nfeats).to(ctx)
score_func.reset_parameters()
score_func(h_src, h_dst, rels).shape == (num_edges)
def tree2(idtype):
"""Generate a tree
1
/ \
4 3
/ \
2 0
Edges are from leaves to root.
"""
g = dgl.DGLGraph().astype(idtype).to(F.ctx())
g.add_nodes(5)
g.add_edge(2, 4)
g.add_edge(0, 4)
g.add_edge(4, 1)
g.add_edge(3, 1)
g.ndata['h'] = F.tensor([0, 1, 2, 3, 4])
g.edata['h'] = F.randn((4, 10))
return g
def test_gin_conv():
g = dgl.DGLGraph(nx.erdos_renyi_graph(20, 0.3))
ctx = F.ctx()
gin_conv = nn.GINConv(lambda x: x, 'mean', 0.1)
gin_conv.initialize(ctx=ctx)
print(gin_conv)
# test #1: basic
feat = F.randn((g.number_of_nodes(), 5))
h = gin_conv(g, feat)
assert h.shape == (20, 5)
# test #2: bipartite
g = dgl.bipartite(sp.sparse.random(100, 200, density=0.1))
feat = (F.randn((100, 5)), F.randn((200, 5)))
h = gin_conv(g, feat)
return h.shape == (20, 5)
def test_glob_att_pool():
ctx = F.ctx()
g = dgl.DGLGraph(nx.path_graph(10))
gap = nn.GlobalAttentionPooling(th.nn.Linear(5, 1), th.nn.Linear(5, 10))
gap = gap.to(ctx)
print(gap)
# test#1: basic
h0 = F.randn((g.number_of_nodes(), 5))
h1 = gap(g, h0)
assert h1.shape[0] == 1 and h1.shape[1] == 10 and h1.dim() == 2
# test#2: batched graph
bg = dgl.batch([g, g, g, g])
h0 = F.randn((bg.number_of_nodes(), 5))
h1 = gap(bg, h0)
assert h1.shape[0] == 4 and h1.shape[1] == 10 and h1.dim() == 2
def test_dense_cheb_conv():
for k in range(1, 4):
ctx = F.ctx()
g = dgl.DGLGraph(sp.sparse.random(100, 100, density=0.1),
readonly=True)
adj = g.adjacency_matrix(ctx=ctx).to_dense()
cheb = nn.ChebConv(5, 2, k)
dense_cheb = nn.DenseChebConv(5, 2, k)
for i in range(len(cheb.fc)):
dense_cheb.W.data[i] = cheb.fc[i].weight.data.t()
if cheb.bias is not None:
dense_cheb.bias.data = cheb.bias.data
feat = F.randn((100, 5))
cheb = cheb.to(ctx)
dense_cheb = dense_cheb.to(ctx)
out_cheb = cheb(g, feat, [2.0])
out_dense_cheb = dense_cheb(adj, feat, 2.0)
assert F.allclose(out_cheb, out_dense_cheb)
def generate_graph(grad=False, add_data=True):
g = dgl.DGLGraph().to(F.ctx())
g.add_nodes(10)
# 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)
if add_data:
ncol = F.randn((10, D))
ecol = F.randn((17, D))
if grad:
ncol = F.attach_grad(ncol)
ecol = F.attach_grad(ecol)
g.ndata['h'] = ncol
g.edata['l'] = ecol
return g
def test_sgc_conv(g, idtype):
ctx = F.ctx()
g = g.astype(idtype).to(ctx)
# not cached
sgc = nn.SGConv(5, 10, 3)
feat = F.randn((g.number_of_nodes(), 5))
sgc = sgc.to(ctx)
h = sgc(g, feat)
assert h.shape[-1] == 10
# cached
sgc = nn.SGConv(5, 10, 3, True)
sgc = sgc.to(ctx)
h_0 = sgc(g, feat)
h_1 = sgc(g, feat + 1)
assert F.allclose(h_0, h_1)
assert h_0.shape[-1] == 10
def test_hypersparse_query():
g = dgl.DGLGraph()
g = g.to(F.ctx())
g.add_nodes(1000001)
g.add_edges([0], [1])
for i in range(10):
assert g.has_node(i)
assert i in g
assert not g.has_node(1000002)
assert g.edge_id(0, 1) == 0
src, dst = g.find_edges([0])
src, dst, eid = g.in_edges(1, form='all')
src, dst, eid = g.out_edges(0, form='all')
src, dst = g.edges()
assert g.in_degree(0) == 0
assert g.in_degree(1) == 1
assert g.out_degree(0) == 1
assert g.out_degree(1) == 0
def tree1(idtype):
"""Generate a tree
0
/ \
1 2
/ \
3 4
Edges are from leaves to root.
"""
g = dgl.graph(([], [])).astype(idtype).to(F.ctx())
g.add_nodes(5)
g.add_edge(3, 1)
g.add_edge(4, 1)
g.add_edge(1, 0)
g.add_edge(2, 0)
g.ndata['h'] = F.tensor([0, 1, 2, 3, 4])
g.edata['h'] = F.randn((4, 10))
return g
def test_dense_sage_conv(g):
ctx = F.ctx()
adj = g.adjacency_matrix(ctx=ctx).to_dense()
sage = nn.SAGEConv(5, 2, 'gcn')
dense_sage = nn.DenseSAGEConv(5, 2)
dense_sage.fc.weight.data = sage.fc_neigh.weight.data
dense_sage.fc.bias.data = sage.fc_neigh.bias.data
if len(g.ntypes) == 2:
feat = (F.randn(
(g.number_of_src_nodes(), 5)), F.randn(
(g.number_of_dst_nodes(), 5)))
else:
feat = F.randn((g.number_of_nodes(), 5))
sage = sage.to(ctx)
dense_sage = dense_sage.to(ctx)
out_sage = sage(g, feat)
out_dense_sage = dense_sage(adj, feat)
assert F.allclose(out_sage, out_dense_sage), g
def test_nn_conv():
ctx = F.ctx()
g = dgl.DGLGraph(sp.sparse.random(100, 100, density=0.1), readonly=True)
edge_func = th.nn.Linear(4, 5 * 10)
nnconv = nn.NNConv(5, 10, edge_func, 'mean')
feat = F.randn((100, 5))
efeat = F.randn((g.number_of_edges(), 4))
nnconv = nnconv.to(ctx)
h = nnconv(g, feat, efeat)
# currently we only do shape check
assert h.shape[-1] == 10
g = dgl.graph(sp.sparse.random(100, 100, density=0.1))
edge_func = th.nn.Linear(4, 5 * 10)
nnconv = nn.NNConv(5, 10, edge_func, 'mean')
feat = F.randn((100, 5))
efeat = F.randn((g.number_of_edges(), 4))
nnconv = nnconv.to(ctx)
h = nnconv(g, feat, efeat)
# currently we only do shape check
assert h.shape[-1] == 10
g = dgl.bipartite(sp.sparse.random(50, 100, density=0.1))
edge_func = th.nn.Linear(4, 5 * 10)
nnconv = nn.NNConv((5, 2), 10, edge_func, 'mean')
feat = F.randn((50, 5))
feat_dst = F.randn((100, 2))
efeat = F.randn((g.number_of_edges(), 4))
nnconv = nnconv.to(ctx)
h = nnconv(g, (feat, feat_dst), efeat)
# currently we only do shape check
assert h.shape[-1] == 10
g = dgl.graph(sp.sparse.random(100, 100, density=0.001))
seed_nodes = th.unique(g.edges()[1])
block = dgl.to_block(g, seed_nodes)
edge_func = th.nn.Linear(4, 5 * 10)
nnconv = nn.NNConv(5, 10, edge_func, 'mean')
feat = F.randn((block.number_of_src_nodes(), 5))
efeat = F.randn((block.number_of_edges(), 4))
nnconv = nnconv.to(ctx)
h = nnconv(block, feat, efeat)
assert h.shape[0] == block.number_of_dst_nodes()
assert h.shape[-1] == 10
def test_weighted_reduce_readout(g, idtype, reducer):
g = g.astype(idtype).to(F.ctx())
g.ndata['h'] = F.randn((g.number_of_nodes(), 3))
g.ndata['w'] = F.randn((g.number_of_nodes(), 1))
g.edata['h'] = F.randn((g.number_of_edges(), 2))
g.edata['w'] = F.randn((g.number_of_edges(), 1))
# Test.1: node readout
x = dgl.readout_nodes(g, 'h', 'w', op=reducer)
# check correctness
subg = dgl.unbatch(g)
subx = []
for sg in subg:
sx = dgl.readout_nodes(sg, 'h', 'w', op=reducer)
subx.append(sx)
assert F.allclose(x, F.cat(subx, dim=0))
x = getattr(dgl, '{}_nodes'.format(reducer))(g, 'h', 'w')
# check correctness
subg = dgl.unbatch(g)
subx = []
for sg in subg:
sx = getattr(dgl, '{}_nodes'.format(reducer))(sg, 'h', 'w')
subx.append(sx)
assert F.allclose(x, F.cat(subx, dim=0))
# Test.2: edge readout
x = dgl.readout_edges(g, 'h', 'w', op=reducer)
# check correctness
subg = dgl.unbatch(g)
subx = []
for sg in subg:
sx = dgl.readout_edges(sg, 'h', 'w', op=reducer)
subx.append(sx)
assert F.allclose(x, F.cat(subx, dim=0))
x = getattr(dgl, '{}_edges'.format(reducer))(g, 'h', 'w')
# check correctness
subg = dgl.unbatch(g)
subx = []
for sg in subg:
sx = getattr(dgl, '{}_edges'.format(reducer))(sg, 'h', 'w')
subx.append(sx)
assert F.allclose(x, F.cat(subx, dim=0))
def test_sage_conv(aggre_type):
ctx = F.ctx()
g = dgl.DGLGraph(sp.sparse.random(100, 100, density=0.1), readonly=True)
sage = nn.SAGEConv(5, 10, aggre_type)
feat = F.randn((100, 5))
h = sage(g, feat)
assert h.shape[-1] == 10
g = dgl.graph(sp.sparse.random(100, 100, density=0.1))
sage = nn.SAGEConv(5, 10, aggre_type)
feat = F.randn((100, 5))
h = sage(g, feat)
assert h.shape[-1] == 10
g = dgl.bipartite(sp.sparse.random(100, 200, density=0.1))
dst_dim = 5 if aggre_type != 'gcn' else 10
sage = nn.SAGEConv((10, dst_dim), 2, aggre_type)
feat = (F.randn((100, 10)), F.randn((200, dst_dim)))
h = sage(g, feat)
assert h.shape[-1] == 2
assert h.shape[0] == 200
g = dgl.graph(sp.sparse.random(100, 100, density=0.001))
seed_nodes = np.unique(g.edges()[1].numpy())
block = dgl.to_block(g, seed_nodes)
sage = nn.SAGEConv(5, 10, aggre_type)
feat = F.randn((block.number_of_src_nodes(), 5))
h = sage(block, feat)
assert h.shape[0] == block.number_of_dst_nodes()
assert h.shape[-1] == 10
# Test the case for graphs without edges
g = dgl.bipartite([], num_nodes=(5, 3))
sage = nn.SAGEConv((3, 3), 2, 'gcn')
feat = (F.randn((5, 3)), F.randn((3, 3)))
h = sage(g, feat)
assert h.shape[-1] == 2
assert h.shape[0] == 3
for aggre_type in ['mean', 'pool', 'lstm']:
sage = nn.SAGEConv((3, 1), 2, aggre_type)
feat = (F.randn((5, 3)), F.randn((3, 1)))
h = sage(g, feat)
assert h.shape[-1] == 2
assert h.shape[0] == 3
def test_prop_edges_dfs(idtype):
g = dgl.graph(nx.path_graph(5), idtype=idtype, device=F.ctx())
g.ndata['x'] = F.ones((5, 2))
dgl.prop_edges_dfs(g, 0, message_func=mfunc, reduce_func=rfunc, apply_node_func=None)
# snr using dfs results in a cumsum
assert F.allclose(g.ndata['x'],
F.tensor([[1., 1.], [2., 2.], [3., 3.], [4., 4.], [5., 5.]]))
g.ndata['x'] = F.ones((5, 2))
dgl.prop_edges_dfs(g, 0, has_reverse_edge=True, message_func=mfunc, reduce_func=rfunc, apply_node_func=None)
# result is cumsum[i] + cumsum[i-1]
assert F.allclose(g.ndata['x'],
F.tensor([[1., 1.], [3., 3.], [5., 5.], [7., 7.], [9., 9.]]))
g.ndata['x'] = F.ones((5, 2))
dgl.prop_edges_dfs(g, 0, has_nontree_edge=True, message_func=mfunc, reduce_func=rfunc, apply_node_func=None)
# result is cumsum[i] + cumsum[i+1]
assert F.allclose(g.ndata['x'],
F.tensor([[3., 3.], [5., 5.], [7., 7.], [9., 9.], [5., 5.]]))
def test_broadcast(idtype, g):
g = g.astype(idtype).to(F.ctx())
gfeat = F.randn((g.batch_size, 3))
# Test.0: broadcast_nodes
g.ndata['h'] = dgl.broadcast_nodes(g, gfeat)
subg = dgl.unbatch(g)
for i, sg in enumerate(subg):
assert F.allclose(
sg.ndata['h'],
F.repeat(F.reshape(gfeat[i], (1, 3)), sg.number_of_nodes(), dim=0))
# Test.1: broadcast_edges
g.edata['h'] = dgl.broadcast_edges(g, gfeat)
subg = dgl.unbatch(g)
for i, sg in enumerate(subg):
assert F.allclose(
sg.edata['h'],
F.repeat(F.reshape(gfeat[i], (1, 3)), sg.number_of_edges(), dim=0))
def test_sage_conv():
for aggre_type in ['mean', 'pool', 'gcn', 'lstm']:
ctx = F.ctx()
g = dgl.DGLGraph(sp.sparse.random(100, 100, density=0.1),
readonly=True)
sage = nn.SAGEConv(5, 10, aggre_type)
feat = F.randn((100, 5))
sage = sage.to(ctx)
h = sage(g, feat)
assert h.shape[-1] == 10
g = dgl.bipartite(sp.sparse.random(100, 200, density=0.1))
dst_dim = 5 if aggre_type != 'gcn' else 10
sage = nn.SAGEConv((10, dst_dim), 2, aggre_type)
feat = (F.randn((100, 10)), F.randn((200, dst_dim)))
sage = sage.to(ctx)
h = sage(g, feat)
assert h.shape[-1] == 2
assert h.shape[0] == 200
def test_dense_sage_conv(g):
ctx = F.ctx()
adj = g.adjacency_matrix(ctx=ctx).tostype('default')
sage = nn.SAGEConv(5, 2, 'gcn')
dense_sage = nn.DenseSAGEConv(5, 2)
sage.initialize(ctx=ctx)
dense_sage.initialize(ctx=ctx)
dense_sage.fc.weight.set_data(sage.fc_neigh.weight.data())
dense_sage.fc.bias.set_data(sage.fc_neigh.bias.data())
if len(g.ntypes) == 2:
feat = (F.randn(
(g.number_of_src_nodes(), 5)), F.randn(
(g.number_of_dst_nodes(), 5)))
else:
feat = F.randn((g.number_of_nodes(), 5))
out_sage = sage(g, feat)
out_dense_sage = dense_sage(adj, feat)
assert F.allclose(out_sage, out_dense_sage)
