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))