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)