def test_tagconv():
g = dgl.DGLGraph(nx.path_graph(3))
ctx = F.ctx()
adj = g.adjacency_matrix(ctx=ctx)
norm = th.pow(g.in_degrees().float(), -0.5)
conv = nn.TAGConv(5, 2, bias=True)
if F.gpu_ctx():
conv = conv.to(ctx)
print(conv)
# test#1: basic
h0 = F.ones((3, 5))
h1 = conv(g, h0)
assert len(g.ndata) == 0
assert len(g.edata) == 0
shp = norm.shape + (1,) * (h0.dim() - 1)
norm = th.reshape(norm, shp).to(ctx)
assert F.allclose(h1, _S2AXWb(adj, norm, h0, conv.lin.weight, conv.lin.bias))
conv = nn.TAGConv(5, 2)
if F.gpu_ctx():
conv = conv.to(ctx)
# test#2: basic
h0 = F.ones((3, 5))
h1 = conv(g, h0)
assert h1.shape[-1] == 2
# test reset_parameters
old_weight = deepcopy(conv.lin.weight.data)
conv.reset_parameters()
new_weight = conv.lin.weight.data
assert not F.allclose(old_weight, new_weight)
def test_simple_pool():
ctx = F.ctx()
g = dgl.DGLGraph(nx.path_graph(15))
sum_pool = nn.SumPooling()
avg_pool = nn.AvgPooling()
max_pool = nn.MaxPooling()
sort_pool = nn.SortPooling(10) # k = 10
print(sum_pool, avg_pool, max_pool, sort_pool)
# test#1: basic
h0 = F.randn((g.number_of_nodes(), 5))
if F.gpu_ctx():
sum_pool = sum_pool.to(ctx)
avg_pool = avg_pool.to(ctx)
max_pool = max_pool.to(ctx)
sort_pool = sort_pool.to(ctx)
h0 = h0.to(ctx)
h1 = sum_pool(g, h0)
assert F.allclose(h1, F.sum(h0, 0))
h1 = avg_pool(g, h0)
assert F.allclose(h1, F.mean(h0, 0))
h1 = max_pool(g, h0)
assert F.allclose(h1, F.max(h0, 0))
h1 = sort_pool(g, h0)
assert h1.shape[0] == 10 * 5 and h1.dim() == 1
# test#2: batched graph
g_ = dgl.DGLGraph(nx.path_graph(5))
bg = dgl.batch([g, g_, g, g_, g])
h0 = F.randn((bg.number_of_nodes(), 5))
if F.gpu_ctx():
h0 = h0.to(ctx)
h1 = sum_pool(bg, h0)
truth = th.stack([F.sum(h0[:15], 0),
F.sum(h0[15:20], 0),
F.sum(h0[20:35], 0),
F.sum(h0[35:40], 0),
F.sum(h0[40:55], 0)], 0)
assert F.allclose(h1, truth)
h1 = avg_pool(bg, h0)
truth = th.stack([F.mean(h0[:15], 0),
F.mean(h0[15:20], 0),
F.mean(h0[20:35], 0),
F.mean(h0[35:40], 0),
F.mean(h0[40:55], 0)], 0)
assert F.allclose(h1, truth)
h1 = max_pool(bg, h0)
truth = th.stack([F.max(h0[:15], 0),
F.max(h0[15:20], 0),
F.max(h0[20:35], 0),
F.max(h0[35:40], 0),
F.max(h0[40:55], 0)], 0)
assert F.allclose(h1, truth)
h1 = sort_pool(bg, h0)
assert h1.shape[0] == 5 and h1.shape[1] == 10 * 5 and h1.dim() == 2
def test_graph_conv():
g = dgl.DGLGraph(nx.path_graph(3))
ctx = F.ctx()
adj = g.adjacency_matrix(ctx=ctx)
conv = nn.GraphConv(5, 2, norm=False, bias=True)
if F.gpu_ctx():
conv = conv.to(ctx)
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)
if F.gpu_ctx():
conv = conv.to(ctx)
# 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)
if F.gpu_ctx():
conv = conv.to(ctx)
# 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
# test rest_parameters
old_weight = deepcopy(conv.weight.data)
conv.reset_parameters()
new_weight = conv.weight.data
assert not F.allclose(old_weight, new_weight)
def test_set_trans():
ctx = F.ctx()
g = dgl.DGLGraph(nx.path_graph(15))
st_enc_0 = nn.SetTransformerEncoder(50, 5, 10, 100, 2, 'sab')
st_enc_1 = nn.SetTransformerEncoder(50, 5, 10, 100, 2, 'isab', 3)
st_dec = nn.SetTransformerDecoder(50, 5, 10, 100, 2, 4)
if F.gpu_ctx():
st_enc_0 = st_enc_0.to(ctx)
st_enc_1 = st_enc_1.to(ctx)
st_dec = st_dec.to(ctx)
print(st_enc_0, st_enc_1, st_dec)
# test#1: basic
h0 = F.randn((g.number_of_nodes(), 50))
h1 = st_enc_0(g, h0)
assert h1.shape == h0.shape
h1 = st_enc_1(g, h0)
assert h1.shape == h0.shape
h2 = st_dec(g, h1)
assert h2.shape[0] == 200 and h2.dim() == 1
# test#2: batched graph
g1 = dgl.DGLGraph(nx.path_graph(5))
g2 = dgl.DGLGraph(nx.path_graph(10))
bg = dgl.batch([g, g1, g2])
h0 = F.randn((bg.number_of_nodes(), 50))
h1 = st_enc_0(bg, h0)
assert h1.shape == h0.shape
h1 = st_enc_1(bg, h0)
assert h1.shape == h0.shape
h2 = st_dec(bg, h1)
assert h2.shape[0] == 3 and h2.shape[1] == 200 and h2.dim() == 2
def test_fps_start_idx():
N = 1000
batch_size = 5
sample_points = 10
x = th.tensor(np.random.uniform(size=(batch_size, int(N / batch_size), 3)))
ctx = F.ctx()
if F.gpu_ctx():
x = x.to(ctx)
res = farthest_point_sampler(x, sample_points, start_idx=0)
assert th.any(res[:, 0] == 0)
def test_agnn_conv():
ctx = F.ctx()
g = dgl.DGLGraph(sp.sparse.random(100, 100, density=0.1), readonly=True)
agnn = nn.AGNNConv(1)
feat = F.randn((100, 5))
if F.gpu_ctx():
agnn = agnn.to(ctx)
h = agnn(g, feat)
assert h.shape[-1] == 5
def test_appnp_conv():
ctx = F.ctx()
g = dgl.DGLGraph(sp.sparse.random(100, 100, density=0.1), readonly=True)
appnp = nn.APPNPConv(10, 0.1)
feat = F.randn((100, 5))
if F.gpu_ctx():
appnp = appnp.to(ctx)
h = appnp(g, feat)
assert h.shape[-1] == 5
def test_gat_conv():
ctx = F.ctx()
g = dgl.DGLGraph(sp.sparse.random(100, 100, density=0.1), readonly=True)
gat = nn.GATConv(5, 2, 4)
feat = F.randn((100, 5))
if F.gpu_ctx():
gat = gat.to(ctx)
h = gat(g, feat)
assert h.shape[-1] == 2 and h.shape[-2] == 4
def test_fps():
N = 1000
batch_size = 5
sample_points = 10
x = th.tensor(np.random.uniform(size=(batch_size, int(N / batch_size), 3)))
ctx = F.ctx()
if F.gpu_ctx():
x = x.to(ctx)
res = farthest_point_sampler(x, sample_points)
assert res.shape[0] == batch_size
assert res.shape[1] == sample_points
assert res.sum() > 0
def test_fps():
N = 1000
batch_size = 5
sample_points = 10
x = mx.nd.array(np.random.uniform(size=(batch_size, int(N/batch_size), 3)))
ctx = F.ctx()
if F.gpu_ctx():
x = x.as_in_context(ctx)
res = farthest_point_sampler(x, sample_points)
assert res.shape[0] == batch_size
assert res.shape[1] == sample_points
assert res.sum() > 0
def test_gin_conv():
for aggregator_type in ['mean', 'max', 'sum']:
ctx = F.ctx()
g = dgl.DGLGraph(sp.sparse.random(100, 100, density=0.1),
readonly=True)
gin = nn.GINConv(th.nn.Linear(5, 12), aggregator_type)
feat = F.randn((100, 5))
if F.gpu_ctx():
gin = gin.to(ctx)
h = gin(g, feat)
assert h.shape[-1] == 12
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))
if F.gpu_ctx():
sage = sage.to(ctx)
h = sage(g, feat)
assert h.shape[-1] == 10
def test_gmm_conv():
ctx = F.ctx()
g = dgl.DGLGraph(sp.sparse.random(100, 100, density=0.1), readonly=True)
gmmconv = nn.GMMConv(5, 10, 3, 4, 'mean')
feat = F.randn((100, 5))
pseudo = F.randn((g.number_of_edges(), 3))
if F.gpu_ctx():
gmmconv = gmmconv.to(ctx)
h = gmmconv(g, feat, pseudo)
# currently we only do shape check
assert h.shape[-1] == 10
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))
if F.gpu_ctx():
sgc = sgc.to(ctx)
h = sgc(g, feat)
assert h.shape[-1] == 10
# cached
sgc = nn.SGConv(5, 10, 3, True)
if F.gpu_ctx():
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_gated_graph_conv():
ctx = F.ctx()
g = dgl.DGLGraph(sp.sparse.random(100, 100, density=0.1), readonly=True)
ggconv = nn.GatedGraphConv(5, 10, 5, 3)
etypes = th.arange(g.number_of_edges()) % 3
feat = F.randn((100, 5))
if F.gpu_ctx():
ggconv = ggconv.to(ctx)
etypes = etypes.to(ctx)
h = ggconv(g, feat, etypes)
# current we only do shape check
assert h.shape[-1] == 10
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))
if F.gpu_ctx():
nnconv = nnconv.to(ctx)
h = nnconv(g, feat, efeat)
# currently we only do shape check
assert h.shape[-1] == 10
def test_cf_conv():
g = dgl.DGLGraph(sp.sparse.random(100, 100, density=0.1), readonly=True)
cfconv = nn.CFConv(node_in_feats=2,
edge_in_feats=3,
hidden_feats=2,
out_feats=3)
ctx = F.ctx()
if F.gpu_ctx():
cfconv = cfconv.to(ctx)
node_feats = F.randn((100, 2))
edge_feats = F.randn((g.number_of_edges(), 3))
h = cfconv(g, node_feats, edge_feats)
# current we only do shape check
assert h.shape[-1] == 3
def test_cf_conv(g, idtype):
g = g.astype(idtype).to(F.ctx())
cfconv = nn.CFConv(node_in_feats=2,
edge_in_feats=3,
hidden_feats=2,
out_feats=3)
ctx = F.ctx()
if F.gpu_ctx():
cfconv = cfconv.to(ctx)
node_feats = F.randn((g.number_of_nodes(), 2))
edge_feats = F.randn((g.number_of_edges(), 3))
h = cfconv(g, node_feats, edge_feats)
# current we only do shape check
assert h.shape[-1] == 3
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)
out_sage = sage(g, feat)
out_dense_sage = dense_sage(adj, feat)
assert F.allclose(out_sage, out_dense_sage)
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)
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_atomic_conv(g, idtype):
g = g.astype(idtype).to(F.ctx())
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((g.number_of_nodes(), 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_glob_att_pool():
g = dgl.DGLGraph(nx.path_graph(10))
gap = nn.GlobalAttentionPooling(th.nn.Linear(5, 1), th.nn.Linear(5, 10))
if F.gpu_ctx():
gap.cuda()
print(gap)
# test#1: basic
h0 = F.randn((g.number_of_nodes(), 5))
h1 = gap(h0, g)
assert h1.shape[0] == 10 and h1.dim() == 1
# test#2: batched graph
bg = dgl.batch([g, g, g, g])
h0 = F.randn((bg.number_of_nodes(), 5))
h1 = gap(h0, bg)
assert h1.shape[0] == 4 and h1.shape[1] == 10 and h1.dim() == 2
def test_set2set():
g = dgl.DGLGraph(nx.path_graph(10))
s2s = nn.Set2Set(5, 3, 3) # hidden size 5, 3 iters, 3 layers
if F.gpu_ctx():
s2s.cuda()
print(s2s)
# test#1: basic
h0 = F.randn((g.number_of_nodes(), 5))
h1 = s2s(h0, g)
assert h1.shape[0] == 10 and h1.dim() == 1
# test#2: batched graph
g1 = dgl.DGLGraph(nx.path_graph(11))
g2 = dgl.DGLGraph(nx.path_graph(5))
bg = dgl.batch([g, g1, g2])
h0 = F.randn((bg.number_of_nodes(), 5))
h1 = s2s(h0, bg)
assert h1.shape[0] == 3 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))
if F.gpu_ctx():
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 test_cf_conv(g, idtype, out_dim):
g = g.astype(idtype).to(F.ctx())
cfconv = nn.CFConv(node_in_feats=2,
edge_in_feats=3,
hidden_feats=2,
out_feats=out_dim)
ctx = F.ctx()
if F.gpu_ctx():
cfconv = cfconv.to(ctx)
src_feats = F.randn((g.number_of_src_nodes(), 2))
edge_feats = F.randn((g.number_of_edges(), 3))
h = cfconv(g, src_feats, edge_feats)
# current we only do shape check
assert h.shape[-1] == out_dim
# case for bipartite graphs
dst_feats = F.randn((g.number_of_dst_nodes(), 3))
h = cfconv(g, (src_feats, dst_feats), edge_feats)
# current we only do shape check
assert h.shape[-1] == out_dim
def test_hetero_conv(agg):
g = dgl.heterograph({
('user', 'follows', 'user'): [(0, 1), (0, 2), (2, 1), (1, 3)],
('user', 'plays', 'game'): [(0, 0), (0, 2), (0, 3), (1, 0), (2, 2)],
('store', 'sells', 'game'): [(0, 0), (0, 3), (1, 1), (1, 2)]})
conv = nn.HeteroGraphConv({
'follows': nn.GraphConv(2, 3),
'plays': nn.GraphConv(2, 4),
'sells': nn.GraphConv(3, 4)},
agg)
if F.gpu_ctx():
conv = conv.to(F.ctx())
uf = F.randn((4, 2))
gf = F.randn((4, 4))
sf = F.randn((2, 3))
uf_dst = F.randn((4, 3))
gf_dst = F.randn((4, 4))
h = conv(g, {'user': uf})
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, 1, 4)
h = conv(g, {'user': uf, '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)
h = conv(g, {'store': sf})
assert set(h.keys()) == {'game'}
if agg != 'stack':
assert h['game'].shape == (4, 4)
else:
assert h['game'].shape == (4, 1, 4)
# test with pair input
conv = nn.HeteroGraphConv({
'follows': nn.SAGEConv(2, 3, 'mean'),
'plays': nn.SAGEConv((2, 4), 4, 'mean'),
'sells': nn.SAGEConv(3, 4, 'mean')},
agg)
if F.gpu_ctx():
conv = conv.to(F.ctx())
h = conv(g, ({'user': uf}, {'user' : uf, '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, 1, 4)
# pair input requires both src and dst type features to be provided
h = conv(g, ({'user': uf}, {'game' : gf}))
assert set(h.keys()) == {'game'}
if agg != 'stack':
assert h['game'].shape == (4, 4)
else:
assert h['game'].shape == (4, 1, 4)
# test with mod args
class MyMod(th.nn.Module):
def __init__(self, s1, s2):
super(MyMod, self).__init__()
self.carg1 = 0
self.carg2 = 0
self.s1 = s1
self.s2 = s2
def forward(self, g, h, arg1=None, *, arg2=None):
if arg1 is not None:
self.carg1 += 1
if arg2 is not None:
self.carg2 += 1
return th.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)
if F.gpu_ctx():
conv = conv.to(F.ctx())
mod_args = {'follows' : (1,), 'plays' : (1,)}
mod_kwargs = {'sells' : {'arg2' : 'abc'}}
h = conv(g, {'user' : uf, '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
# Test multiple nodes
sg, inv = dgl.khop_out_subgraph(g, {
'user': F.tensor([2], idtype),
'game': 0
},
k=1)
assert sg.num_edges('follows') == 0
u, v = sg['plays'].edges()
edge_set = set(zip(list(F.asnumpy(u)), list(F.asnumpy(v))))
assert edge_set == {(0, 1)}
assert F.array_equal(F.astype(inv['user'], idtype), F.tensor([0], idtype))
assert F.array_equal(F.astype(inv['game'], idtype), F.tensor([0], idtype))
@unittest.skipIf(not F.gpu_ctx(), 'only necessary with GPU')
@unittest.skipIf(dgl.backend.backend_name != "pytorch",
reason="UVA only supported for PyTorch")
@pytest.mark.parametrize('parent_idx_device', [('cpu', F.cpu()),
('cuda', F.cuda()),
('uva', F.cpu()),
('uva', F.cuda())])
@pytest.mark.parametrize('child_device', [F.cpu(), F.cuda()])
def test_subframes(parent_idx_device, child_device):
parent_device, idx_device = parent_idx_device
g = dgl.graph((F.tensor([1, 2, 3],
dtype=F.int64), F.tensor([2, 3, 4],
dtype=F.int64)))
print(g.device)
g.ndata['x'] = F.randn((5, 4))
g.edata['a'] = F.randn((3, 6))
