def test_sddmm(g, shp, lhs_target, rhs_target, msg, index_dtype): if dgl.backend.backend_name == 'mxnet' and g.number_of_edges() == 0: pytest.skip() # mxnet do not support zero shape tensor if dgl.backend.backend_name == 'tensorflow' and index_dtype == 'int32': pytest.skip() # tensorflow dlpack has problem with int32 ndarray. if index_dtype == 'int32': g = g.int() else: g = g.long() print(g) print(g.idtype) len_lhs = select(lhs_target, g.number_of_src_nodes(), g.number_of_edges(), g.number_of_dst_nodes()) lhs_shp = (len_lhs, ) + shp[0] len_rhs = select(rhs_target, g.number_of_src_nodes(), g.number_of_edges(), g.number_of_dst_nodes()) rhs_shp = (len_rhs, ) + shp[1] feat_lhs = F.tensor(np.random.rand(*lhs_shp) + 1) feat_rhs = F.tensor(np.random.rand(*rhs_shp) + 1) print('lhs shape: {}, rhs shape: {}'.format(F.shape(feat_lhs), F.shape(feat_rhs))) lhs_frame = select(lhs_target, g.srcdata, g.edata, g.dstdata) rhs_frame = select(rhs_target, g.srcdata, g.edata, g.dstdata) lhs_frame['x'] = F.attach_grad(F.clone(feat_lhs)) rhs_frame['y'] = F.attach_grad(F.clone(feat_rhs)) msg_func = lhs_target + '_' + msg + '_' + rhs_target print('SDDMM(message func: {})'.format(msg_func)) lhs = F.attach_grad(F.clone(feat_lhs)) rhs = F.attach_grad(F.clone(feat_rhs)) with F.record_grad(): e = gsddmm(g, msg, lhs, rhs, lhs_target=lhs_target, rhs_target=rhs_target) F.backward(F.reduce_sum(e)) grad_lhs = F.grad(lhs) grad_rhs = F.grad(rhs) with F.record_grad(): g.apply_edges(udf_apply_edges[msg_func]) if g.number_of_edges() > 0: e1 = g.edata['m'] assert F.allclose(e, e1) print('forward passed') F.backward(F.reduce_sum(e1)) if msg != 'copy_rhs': assert F.allclose(F.grad(lhs_frame['x']), grad_lhs) if msg != 'copy_lhs': assert F.allclose(F.grad(rhs_frame['y']), grad_rhs) print('backward passed') lhs_frame.pop('x') rhs_frame.pop('y') if 'm' in g.edata: g.edata.pop('m')
def test_segment_reduce(reducer): ctx = F.ctx() value = F.tensor(np.random.rand(10, 5)) v1 = F.attach_grad(F.clone(value)) v2 = F.attach_grad(F.clone(value)) seglen = F.tensor([2, 3, 0, 4, 1, 0, 0]) u = F.copy_to(F.arange(0, F.shape(value)[0], F.int32), ctx) v = F.repeat(F.copy_to(F.arange(0, len(seglen), F.int32), ctx), seglen, dim=0) num_nodes = {'_U': len(u), '_V': len(seglen)} g = dgl.convert.heterograph({('_U', '_E', '_V'): (u, v)}, num_nodes_dict=num_nodes) with F.record_grad(): rst1 = gspmm(g, 'copy_lhs', reducer, v1, None) if reducer in ['max', 'min']: rst1 = F.replace_inf_with_zero(rst1) F.backward(F.reduce_sum(rst1)) grad1 = F.grad(v1) with F.record_grad(): rst2 = segment_reduce(seglen, v2, reducer=reducer) F.backward(F.reduce_sum(rst2)) assert F.allclose(rst1, rst2) print('forward passed') grad2 = F.grad(v2) assert F.allclose(grad1, grad2) print('backward passed')
def _test(apply_func): g = generate_graph() f = g.ndata['f'] # one out place run to get result g.send((u, v), message_func) g.recv([0, 1, 2, 3, 9], reduce_func, apply_func) result = g.get_n_repr()['f'] # inplace deg bucket run v1 = F.clone(f) g.ndata['f'] = v1 g.send((u, v), message_func) g.recv([0, 1, 2, 3, 9], reduce_func, apply_func, inplace=True) r1 = g.get_n_repr()['f'] # check result assert F.allclose(r1, result) # check inplace assert F.allclose(v1, r1) # inplace e2v v1 = F.clone(f) g.ndata['f'] = v1 g.send((u, v), message_func) g.recv([0, 1, 2, 3, 9], fn.sum(msg='m', out='f'), apply_func, inplace=True) r1 = g.ndata['f'] # check result assert F.allclose(r1, result) # check inplace assert F.allclose(v1, r1)
def test_edge_softmax(g, norm_by, shp, idtype): g = g.astype(idtype).to(F.ctx()) edata = F.tensor(np.random.rand(g.number_of_edges(), *shp)) e1 = F.attach_grad(F.clone(edata)) with F.record_grad(): score1 = edge_softmax(g, e1, norm_by=norm_by) F.backward(F.reduce_sum(score1)) grad_edata = F.grad(e1) with F.record_grad(): e2 = F.attach_grad(F.clone(edata)) e2_2d = F.reshape( e2, (g.number_of_src_nodes(), g.number_of_dst_nodes(), *e2.shape[1:])) if norm_by == 'src': score2 = F.softmax(e2_2d, 1) score2 = F.reshape(score2, (-1, *e2.shape[1:])) if norm_by == 'dst': score2 = F.softmax(e2_2d, 0) score2 = F.reshape(score2, (-1, *e2.shape[1:])) assert F.allclose(score1, score2) print('forward passed') F.backward(F.reduce_sum(score2)) assert F.allclose(F.grad(e2), grad_edata) print('backward passed')
def test_spmm(idtype, g, shp, msg, reducer): g = g.astype(idtype).to(F.ctx()) if dgl.backend.backend_name == 'tensorflow' and (reducer in ['min', 'max']): pytest.skip() # tensorflow dlpack has problem writing into int32 arrays on GPU. print(g) print(g.idtype) hu = F.tensor(np.random.rand(*((g.number_of_src_nodes(),) + shp[0])) + 1) he = F.tensor(np.random.rand(*((g.number_of_edges(),) + shp[1])) + 1) print('u shape: {}, e shape: {}'.format(F.shape(hu), F.shape(he))) g.srcdata['x'] = F.attach_grad(F.clone(hu)) g.edata['w'] = F.attach_grad(F.clone(he)) print('SpMM(message func: {}, reduce func: {})'.format(msg, reducer)) u = F.attach_grad(F.clone(hu)) e = F.attach_grad(F.clone(he)) with F.record_grad(): v = gspmm(g, msg, reducer, u, e) non_degree_indices = F.tensor( np.nonzero(F.asnumpy(g.in_degrees()) != 0)[0]) v = F.gather_row(v, non_degree_indices) if g.number_of_edges() > 0: F.backward(F.reduce_sum(v)) if msg != 'copy_rhs': grad_u = F.grad(u) if msg != 'copy_lhs': grad_e = F.grad(e) with F.record_grad(): g.update_all(udf_msg[msg], udf_reduce[reducer]) if g.number_of_edges() > 0: v1 = F.gather_row(g.dstdata['v'], non_degree_indices) assert F.allclose(v, v1) print('forward passed') F.backward(F.reduce_sum(v1)) if msg != 'copy_rhs': if reducer in ['min', 'max']: # there might be some numerical errors rate = F.reduce_sum(F.abs(F.grad(g.srcdata['x']) - grad_u)) /\ F.reduce_sum(F.abs(grad_u)) assert F.as_scalar(rate) < 1e-2, rate else: assert F.allclose(F.grad(g.srcdata['x']), grad_u) if msg != 'copy_lhs': if reducer in ['min', 'max']: rate = F.reduce_sum(F.abs(F.grad(g.edata['w']) - grad_e)) /\ F.reduce_sum(F.abs(grad_e)) assert F.as_scalar(rate) < 1e-2, rate else: assert F.allclose(F.grad(g.edata['w']), grad_e) print('backward passed') g.srcdata.pop('x') g.edata.pop('w') if 'v' in g.dstdata: g.dstdata.pop('v')
def test_spmm(idtype, g, shp, msg, reducer): g = g.astype(idtype).to(F.ctx()) print(g) print(g.idtype) hu = F.tensor(np.random.rand(*((g.number_of_src_nodes(), ) + shp[0])) + 1) he = F.tensor(np.random.rand(*((g.number_of_edges(), ) + shp[1])) + 1) print('u shape: {}, e shape: {}'.format(F.shape(hu), F.shape(he))) g.srcdata['x'] = F.attach_grad(F.clone(hu)) g.edata['w'] = F.attach_grad(F.clone(he)) print('SpMM(message func: {}, reduce func: {})'.format(msg, reducer)) u = F.attach_grad(F.clone(hu)) e = F.attach_grad(F.clone(he)) with F.record_grad(): v = gspmm(g, msg, reducer, u, e) if reducer in ['max', 'min']: v = F.replace_inf_with_zero(v) if g.number_of_edges() > 0: F.backward(F.reduce_sum(v)) if msg != 'copy_rhs': grad_u = F.grad(u) if msg != 'copy_lhs': grad_e = F.grad(e) with F.record_grad(): g.update_all(udf_msg[msg], udf_reduce[reducer]) if g.number_of_edges() > 0: v1 = g.dstdata['v'] assert F.allclose(v, v1) print('forward passed') F.backward(F.reduce_sum(v1)) if msg != 'copy_rhs': if reducer in ['min', 'max']: # there might be some numerical errors rate = F.reduce_sum(F.abs(F.grad(g.srcdata['x']) - grad_u)) /\ F.reduce_sum(F.abs(grad_u)) assert F.as_scalar(rate) < 1e-2, rate else: assert F.allclose(F.grad(g.srcdata['x']), grad_u) if msg != 'copy_lhs': if reducer in ['min', 'max']: rate = F.reduce_sum(F.abs(F.grad(g.edata['w']) - grad_e)) /\ F.reduce_sum(F.abs(grad_e)) assert F.as_scalar(rate) < 1e-2, rate else: assert F.allclose(F.grad(g.edata['w']), grad_e) print('backward passed') g.srcdata.pop('x') g.edata.pop('w') if 'v' in g.dstdata: g.dstdata.pop('v')
def _test(apply_func): g = generate_graph() f = g.ndata['f'] # an out place run to get result g.send_and_recv((u, v), fn.copy_src(src='f', out='m'), fn.sum(msg='m', out='f'), apply_func) result = g.ndata['f'] # inplace deg bucket v1 = F.clone(f) g.ndata['f'] = v1 g.send_and_recv((u, v), message_func, reduce_func, apply_func, inplace=True) r1 = g.ndata['f'] # check result assert F.allclose(r1, result) # check inplace assert F.allclose(v1, r1) # inplace v2v spmv v1 = F.clone(f) g.ndata['f'] = v1 g.send_and_recv((u, v), fn.copy_src(src='f', out='m'), fn.sum(msg='m', out='f'), apply_func, inplace=True) r1 = g.ndata['f'] # check result assert F.allclose(r1, result) # check inplace assert F.allclose(v1, r1) # inplace e2v spmv v1 = F.clone(f) g.ndata['f'] = v1 g.send_and_recv((u, v), message_func, fn.sum(msg='m', out='f'), apply_func, inplace=True) r1 = g.ndata['f'] # check result assert F.allclose(r1, result) # check inplace assert F.allclose(v1, r1)
def _test(red): g = dgl.DGLGraph(nx.erdos_renyi_graph(100, 0.1)) hu, hv, he = generate_feature(g, 'none') g.ndata['u'] = F.attach_grad(F.clone(hu)) g.ndata['v'] = F.attach_grad(F.clone(hv)) g.edata['e'] = F.attach_grad(F.clone(he)) with F.record_grad(): g.update_all(fn.copy_edge(edge='e', out='m'), builtin[red](msg='m', out='r1')) r1 = g.ndata['r1'] F.backward(r1.sum()) e_grad1 = F.grad(g.edata['e']) # reset grad g.ndata['u'] = F.attach_grad(F.clone(hu)) g.ndata['v'] = F.attach_grad(F.clone(hv)) g.edata['e'] = F.attach_grad(F.clone(he)) with F.record_grad(): g.update_all(udf_copy_edge, udf_reduce[red]) r2 = g.ndata['r2'] F.backward(r2.sum()) e_grad2 = F.grad(g.edata['e']) assert F.allclose(r1, r2) assert(F.allclose(e_grad1, e_grad2))
def _test(red, partial): g = dgl.DGLGraph(nx.erdos_renyi_graph(100, 0.1)) # NOTE(zihao): add self-loop to avoid zero-degree nodes. # https://github.com/dmlc/dgl/issues/761 g.add_edges(g.nodes(), g.nodes()) g = g.to(F.ctx()) hu, hv, he = generate_feature(g, 'none', 'none') if partial: nid = F.tensor(list(range(0, 100, 2)), g.idtype) g.ndata['u'] = F.attach_grad(F.clone(hu)) g.ndata['v'] = F.attach_grad(F.clone(hv)) g.edata['e'] = F.attach_grad(F.clone(he)) with F.record_grad(): if partial: g.pull(nid, fn.copy_src(src='u', out='m'), builtin[red](msg='m', out='r1')) else: g.update_all(fn.copy_src(src='u', out='m'), builtin[red](msg='m', out='r1')) r1 = g.ndata['r1'] F.backward(F.reduce_sum(r1)) n_grad1 = F.grad(g.ndata['u']) # reset grad g.ndata['u'] = F.attach_grad(F.clone(hu)) g.ndata['v'] = F.attach_grad(F.clone(hv)) g.edata['e'] = F.attach_grad(F.clone(he)) with F.record_grad(): if partial: g.pull(nid, udf_copy_src, udf_reduce[red]) else: g.update_all(udf_copy_src, udf_reduce[red]) r2 = g.ndata['r2'] F.backward(F.reduce_sum(r2)) n_grad2 = F.grad(g.ndata['u']) def _print_error(a, b): print("ERROR: Test copy_src_{} partial: {}". format(red, partial)) for i, (x, y) in enumerate(zip(F.asnumpy(a).flatten(), F.asnumpy(b).flatten())): if not np.allclose(x, y): print('@{} {} v.s. {}'.format(i, x, y)) if not F.allclose(r1, r2): _print_error(r1, r2) assert F.allclose(r1, r2) if F.backend_name != "jax": if not F.allclose(n_grad1, n_grad2): print('node grad') _print_error(n_grad1, n_grad2) assert(F.allclose(n_grad1, n_grad2))
def _test(red, partial): g = dgl.DGLGraph(nx.erdos_renyi_graph(100, 0.1)) # NOTE(zihao): add self-loop to avoid zero-degree nodes. g.add_edges(g.nodes(), g.nodes()) hu, hv, he = generate_feature(g, 'none', 'none') if partial: nid = F.tensor(list(range(0, 100, 2))) g.ndata['u'] = F.attach_grad(F.clone(hu)) g.ndata['v'] = F.attach_grad(F.clone(hv)) g.edata['e'] = F.attach_grad(F.clone(he)) with F.record_grad(): if partial: g.pull(nid, fn.copy_edge(edge='e', out='m'), builtin[red](msg='m', out='r1')) else: g.update_all(fn.copy_edge(edge='e', out='m'), builtin[red](msg='m', out='r1')) r1 = g.ndata['r1'] F.backward(r1.sum()) e_grad1 = F.grad(g.edata['e']) # reset grad g.ndata['u'] = F.attach_grad(F.clone(hu)) g.ndata['v'] = F.attach_grad(F.clone(hv)) g.edata['e'] = F.attach_grad(F.clone(he)) with F.record_grad(): if partial: g.pull(nid, udf_copy_edge, udf_reduce[red]) else: g.update_all(udf_copy_edge, udf_reduce[red]) r2 = g.ndata['r2'] F.backward(r2.sum()) e_grad2 = F.grad(g.edata['e']) def _print_error(a, b): print("ERROR: Test copy_edge_{} partial: {}".format(red, partial)) for i, (x, y) in enumerate( zip(F.asnumpy(a).flatten(), F.asnumpy(b).flatten())): if not np.allclose(x, y): print('@{} {} v.s. {}'.format(i, x, y)) if not F.allclose(r1, r2): _print_error(r1, r2) assert F.allclose(r1, r2) if not F.allclose(e_grad1, e_grad2): print('edge gradient') _print_error(e_grad1, e_grad2) assert (F.allclose(e_grad1, e_grad2))
def _test(red, partial): g = dgl.DGLGraph(nx.erdos_renyi_graph(100, 0.1)) # NOTE(zihao): add self-loop to avoid zero-degree nodes. # https://github.com/dmlc/dgl/issues/761 g.add_edges(g.nodes(), g.nodes()) hu, hv, he = generate_feature(g, 'none') if partial: nid = F.tensor(list(range(0, 100, 2))) g.ndata['u'] = F.attach_grad(F.clone(hu)) g.ndata['v'] = F.attach_grad(F.clone(hv)) g.edata['e'] = F.attach_grad(F.clone(he)) with F.record_grad(): if partial: g.pull(nid, fn.copy_src(src='u', out='m'), builtin[red](msg='m', out='r1')) else: g.update_all(fn.copy_src(src='u', out='m'), builtin[red](msg='m', out='r1')) r1 = g.ndata['r1'] F.backward(r1.sum()) n_grad1 = F.grad(g.ndata['u']) # reset grad g.ndata['u'] = F.attach_grad(F.clone(hu)) g.ndata['v'] = F.attach_grad(F.clone(hv)) g.edata['e'] = F.attach_grad(F.clone(he)) with F.record_grad(): if partial: g.pull(nid, udf_copy_src, udf_reduce[red]) else: g.update_all(udf_copy_src, udf_reduce[red]) r2 = g.ndata['r2'] F.backward(r2.sum()) n_grad2 = F.grad(g.ndata['u']) assert F.allclose(r1, r2) assert(F.allclose(n_grad1, n_grad2))
def _test(red, partial): g = dgl.DGLGraph(nx.erdos_renyi_graph(100, 0.1)) hu, hv, he = generate_feature(g, 'none') if partial: nid = F.tensor(list(range(0, 100, 2))) g.ndata['u'] = F.attach_grad(F.clone(hu)) g.ndata['v'] = F.attach_grad(F.clone(hv)) g.edata['e'] = F.attach_grad(F.clone(he)) with F.record_grad(): if partial: g.pull(nid, fn.copy_src(src='u', out='m'), builtin[red](msg='m', out='r1')) else: g.update_all(fn.copy_src(src='u', out='m'), builtin[red](msg='m', out='r1')) r1 = g.ndata['r1'] F.backward(r1.sum()) n_grad1 = F.grad(g.ndata['u']) # reset grad g.ndata['u'] = F.attach_grad(F.clone(hu)) g.ndata['v'] = F.attach_grad(F.clone(hv)) g.edata['e'] = F.attach_grad(F.clone(he)) with F.record_grad(): if partial: g.pull(nid, udf_copy_src, udf_reduce[red]) else: g.update_all(udf_copy_src, udf_reduce[red]) r2 = g.ndata['r2'] F.backward(r2.sum()) n_grad2 = F.grad(g.ndata['u']) assert F.allclose(r1, r2) assert (F.allclose(n_grad1, n_grad2))
def test_copy(): num_layers = 2 g = generate_rand_graph(100) g.ndata['h'] = g.ndata['h1'] nf = create_mini_batch(g, num_layers) nf.copy_from_parent() for i in range(nf.num_layers): assert len(g.ndata.keys()) == len(nf.layers[i].data.keys()) for key in g.ndata.keys(): assert key in nf.layers[i].data.keys() assert_array_equal( F.asnumpy(nf.layers[i].data[key]), F.asnumpy(g.nodes[nf.layer_parent_nid(i)].data[key])) for i in range(nf.num_blocks): assert len(g.edata.keys()) == len(nf.blocks[i].data.keys()) for key in g.edata.keys(): assert key in nf.blocks[i].data.keys() assert_array_equal( F.asnumpy(nf.blocks[i].data[key]), F.asnumpy(g.edges[nf.block_parent_eid(i)].data[key])) nf = create_mini_batch(g, num_layers) node_embed_names = [['h'], ['h1'], ['h']] edge_embed_names = [['h2'], ['h2']] nf.copy_from_parent(node_embed_names=node_embed_names, edge_embed_names=edge_embed_names) for i in range(nf.num_layers): assert len(node_embed_names[i]) == len(nf.layers[i].data.keys()) for key in node_embed_names[i]: assert key in nf.layers[i].data.keys() assert_array_equal( F.asnumpy(nf.layers[i].data[key]), F.asnumpy(g.nodes[nf.layer_parent_nid(i)].data[key])) for i in range(nf.num_blocks): assert len(edge_embed_names[i]) == len(nf.blocks[i].data.keys()) for key in edge_embed_names[i]: assert key in nf.blocks[i].data.keys() assert_array_equal( F.asnumpy(nf.blocks[i].data[key]), F.asnumpy(g.edges[nf.block_parent_eid(i)].data[key])) nf = create_mini_batch(g, num_layers) g.ndata['h0'] = F.clone(g.ndata['h']) node_embed_names = [['h0'], [], []] nf.copy_from_parent(node_embed_names=node_embed_names, edge_embed_names=None) for i in range(num_layers): nf.block_compute(i, fn.copy_src(src='h%d' % i, out='m'), fn.sum(msg='m', out='t'), lambda nodes: {'h%d' % (i + 1): nodes.data['t'] + 1}) g.update_all(fn.copy_src(src='h', out='m'), fn.sum(msg='m', out='t'), lambda nodes: {'h': nodes.data['t'] + 1}) assert_allclose(F.asnumpy(nf.layers[i + 1].data['h%d' % (i + 1)]), F.asnumpy(g.nodes[nf.layer_parent_nid(i + 1)].data['h']), rtol=1e-4, atol=1e-4) nf.copy_to_parent(node_embed_names=[['h0'], ['h1'], ['h2']]) for i in range(num_layers + 1): assert_array_equal( F.asnumpy(nf.layers[i].data['h%d' % i]), F.asnumpy(g.nodes[nf.layer_parent_nid(i)].data['h%d' % i])) nf = create_mini_batch(g, num_layers) g.ndata['h0'] = F.clone(g.ndata['h']) g.ndata['h1'] = F.clone(g.ndata['h']) g.ndata['h2'] = F.clone(g.ndata['h']) node_embed_names = [['h0'], ['h1'], ['h2']] nf.copy_from_parent(node_embed_names=node_embed_names, edge_embed_names=None) def msg_func(edge, ind): assert 'h%d' % ind in edge.src.keys() return {'m': edge.src['h%d' % ind]} def reduce_func(node, ind): assert 'h%d' % (ind + 1) in node.data.keys() return { 'h': F.sum(node.mailbox['m'], 1) + node.data['h%d' % (ind + 1)] } for i in range(num_layers): nf.block_compute(i, partial(msg_func, ind=i), partial(reduce_func, ind=i))
def _test(g, lhs, rhs, binary_op, reducer, paritial, nid, broadcast='none'): hu, hv, he = generate_feature(g, broadcast) g.ndata['u'] = F.attach_grad(F.clone(hu)) g.ndata['v'] = F.attach_grad(F.clone(hv)) g.edata['e'] = F.attach_grad(F.clone(he)) builtin_msg_name = "{}_{}_{}".format(lhs, binary_op, rhs) builtin_msg = getattr(fn, builtin_msg_name) builtin_red = getattr(fn, reducer) def target_feature_switch(g, target): if target == "u": return g.ndata["u"] elif target == "v": return g.ndata["v"] else: return g.edata["e"] with F.record_grad(): if partial: g.pull(nid, builtin_msg(lhs, rhs, 'm'), builtin_red('m', 'r1')) else: g.update_all(builtin_msg(lhs, rhs, 'm'), builtin_red('m', 'r1')) r1 = g.ndata.pop('r1') F.backward(r1.sum()) lhs_grad_1 = F.grad(target_feature_switch(g, lhs)) rhs_grad_1 = F.grad(target_feature_switch(g, rhs)) # reset grad g.ndata['u'] = F.attach_grad(F.clone(hu)) g.ndata['v'] = F.attach_grad(F.clone(hv)) g.edata['e'] = F.attach_grad(F.clone(he)) def target_switch(edges, target): if target == "u": return edges.src elif target == "v": return edges.dst elif target == "e": return edges.data else: assert (0), "Unknown target {}".format(target) def mfunc(edges): op = getattr(F, binary_op) lhs_data = target_switch(edges, lhs) rhs_data = target_switch(edges, rhs) return {"m": op(lhs_data[lhs], rhs_data[rhs])} def rfunc(nodes): op = getattr(F, reducer) return {"r2": op(nodes.mailbox['m'], 1)} with F.record_grad(): if partial: g.pull(nid, mfunc, rfunc) else: g.update_all(mfunc, rfunc) r2 = g.ndata.pop('r2') F.backward(r2.sum(), F.tensor([1.])) lhs_grad_2 = F.grad(target_feature_switch(g, lhs)) rhs_grad_2 = F.grad(target_feature_switch(g, rhs)) if reducer == 'prod': rtol = 1e-2 atol = 1e-2 else: rtol = 1e-4 atol = 1e-4 def _print_error(a, b): print("ERROR: Test {}_{}_{}_{} {}".format(lhs, binary_op, rhs, reducer, broadcast)) print(a, b) for i, (x, y) in enumerate( zip( F.asnumpy(F.cpu(a)).flatten(), F.asnumpy(F.cpu(b)).flatten())): if not np.allclose(x, y, rtol, atol): print('@{} {} v.s. {}'.format(i, x, y)) if not F.allclose(r1, r2, rtol, atol): _print_error(r1, r2) assert F.allclose(r1, r2, rtol, atol) if not F.allclose(lhs_grad_1, lhs_grad_2, rtol, atol): print("left grad") _print_error(lhs_grad_1, lhs_grad_2) assert (F.allclose(lhs_grad_1, lhs_grad_2, rtol, atol)) if not F.allclose(rhs_grad_1, rhs_grad_2, rtol, atol): print("right grad") _print_error(rhs_grad_1, rhs_grad_2) assert (F.allclose(rhs_grad_1, rhs_grad_2, rtol, atol))
def _test(g, lhs, rhs, binary_op, reducer, partial, nid, broadcast='none'): # initialize node/edge features with uniform(-1, 1) hu, hv, he = generate_feature(g, broadcast, binary_op) if binary_op == 'div': # op = div # lhs range: [-1, 1] # rhs range: [1, 2] # result range: [-1, 1] if rhs == 'u': hu = (hu + 3) / 2 elif rhs == 'v': hv = (hv + 3) / 2 elif rhs == 'e': he = (he + 3) / 2 if binary_op == 'add' or binary_op == 'sub': # op = add, sub # lhs range: [-1/2, 1/2] # rhs range: [-1/2, 1/2] # result range: [-1, 1] hu = hu / 2 hv = hv / 2 he = he / 2 g.ndata['u'] = F.attach_grad(F.clone(hu)) g.ndata['v'] = F.attach_grad(F.clone(hv)) g.edata['e'] = F.attach_grad(F.clone(he)) builtin_msg_name = "{}_{}_{}".format(lhs, binary_op, rhs) builtin_msg = getattr(fn, builtin_msg_name) builtin_red = getattr(fn, reducer) def target_feature_switch(g, target): if target == "u": return g.ndata["u"] elif target == "v": return g.ndata["v"] else: return g.edata["e"] with F.record_grad(): if partial: g.pull(nid, builtin_msg(lhs, rhs, 'm'), builtin_red('m', 'r1')) else: g.update_all(builtin_msg(lhs, rhs, 'm'), builtin_red('m', 'r1')) r1 = g.ndata.pop('r1') F.backward(F.reduce_sum(r1)) lhs_grad_1 = F.grad(target_feature_switch(g, lhs)) rhs_grad_1 = F.grad(target_feature_switch(g, rhs)) # reset grad g.ndata['u'] = F.attach_grad(F.clone(hu)) g.ndata['v'] = F.attach_grad(F.clone(hv)) g.edata['e'] = F.attach_grad(F.clone(he)) def target_switch(edges, target): if target == "u": return edges.src elif target == "v": return edges.dst elif target == "e": return edges.data else: assert(0), "Unknown target {}".format(target) def mfunc(edges): op = getattr(F, binary_op) lhs_data = target_switch(edges, lhs)[lhs] rhs_data = target_switch(edges, rhs)[rhs] # NOTE(zihao): we need to do batched broadcast # e.g. (68, 3, 1) op (68, 5, 3, 4) while F.ndim(lhs_data) < F.ndim(rhs_data): lhs_data = F.unsqueeze(lhs_data, 1) while F.ndim(rhs_data) < F.ndim(lhs_data): rhs_data = F.unsqueeze(rhs_data, 1) return {"m": op(lhs_data, rhs_data)} def rfunc(nodes): op = getattr(F, reducer) return {"r2": op(nodes.mailbox['m'], 1)} with F.record_grad(): if partial: g.pull(nid, mfunc, rfunc) else: g.update_all(mfunc, rfunc) r2 = g.ndata.pop('r2') F.backward(F.reduce_sum(r2), F.tensor([1.])) lhs_grad_2 = F.grad(target_feature_switch(g, lhs)) rhs_grad_2 = F.grad(target_feature_switch(g, rhs)) rtol = 1e-4 atol = 1e-4 def _print_error(a, b): print("ERROR: Test {}_{}_{}_{} broadcast: {} partial: {}". format(lhs, binary_op, rhs, reducer, broadcast, partial)) return if lhs == 'u': lhs_data = hu elif lhs == 'v': lhs_data = hv elif lhs == 'e': lhs_data = he if rhs == 'u': rhs_data = hu elif rhs == 'v': rhs_data = hv elif rhs == 'e': rhs_data = he print("lhs", F.asnumpy(lhs_data).tolist()) print("rhs", F.asnumpy(rhs_data).tolist()) for i, (x, y) in enumerate(zip(F.asnumpy(a).flatten(), F.asnumpy(b).flatten())): if not np.allclose(x, y, rtol, atol): print('@{} {} v.s. {}'.format(i, x, y)) if not F.allclose(r1, r2, rtol, atol): _print_error(r1, r2) assert F.allclose(r1, r2, rtol, atol) if not F.allclose(lhs_grad_1, lhs_grad_2, rtol, atol): print("left grad") _print_error(lhs_grad_1, lhs_grad_2) assert(F.allclose(lhs_grad_1, lhs_grad_2, rtol, atol)) if not F.allclose(rhs_grad_1, rhs_grad_2, rtol, atol): print("right grad") _print_error(rhs_grad_1, rhs_grad_2) assert(F.allclose(rhs_grad_1, rhs_grad_2, rtol, atol))
def test_edge_softmax(g, norm_by, idtype): print("params", norm_by, idtype) g = create_test_heterograph(idtype) x1 = F.randn((g.num_edges('plays'), feat_size)) x2 = F.randn((g.num_edges('follows'), feat_size)) x3 = F.randn((g.num_edges('develops'), feat_size)) x4 = F.randn((g.num_edges('wishes'), feat_size)) F.attach_grad(F.clone(x1)) F.attach_grad(F.clone(x2)) F.attach_grad(F.clone(x3)) F.attach_grad(F.clone(x4)) g['plays'].edata['eid'] = x1 g['follows'].edata['eid'] = x2 g['develops'].edata['eid'] = x3 g['wishes'].edata['eid'] = x4 ################################################################# # edge_softmax() on homogeneous graph ################################################################# with F.record_grad(): hm_g = dgl.to_homogeneous(g) hm_x = F.cat((x3, x2, x1, x4), 0) hm_e = F.attach_grad(F.clone(hm_x)) score_hm = edge_softmax(hm_g, hm_e, norm_by=norm_by) hm_g.edata['score'] = score_hm ht_g = dgl.to_heterogeneous(hm_g, g.ntypes, g.etypes) r1 = ht_g.edata['score'][('user', 'plays', 'game')] r2 = ht_g.edata['score'][('user', 'follows', 'user')] r3 = ht_g.edata['score'][('developer', 'develops', 'game')] r4 = ht_g.edata['score'][('user', 'wishes', 'game')] F.backward(F.reduce_sum(r1) + F.reduce_sum(r2)) grad_edata_hm = F.grad(hm_e) ################################################################# # edge_softmax() on heterogeneous graph ################################################################# e1 = F.attach_grad(F.clone(x1)) e2 = F.attach_grad(F.clone(x2)) e3 = F.attach_grad(F.clone(x3)) e4 = F.attach_grad(F.clone(x4)) e = { ('user', 'follows', 'user'): e2, ('user', 'plays', 'game'): e1, ('user', 'wishes', 'game'): e4, ('developer', 'develops', 'game'): e3 } with F.record_grad(): score = edge_softmax(g, e, norm_by=norm_by) r5 = score[('user', 'plays', 'game')] r6 = score[('user', 'follows', 'user')] r7 = score[('developer', 'develops', 'game')] r8 = score[('user', 'wishes', 'game')] F.backward(F.reduce_sum(r5) + F.reduce_sum(r6)) grad_edata_ht = F.cat((F.grad(e3), F.grad(e2), F.grad(e1), F.grad(e4)), 0) # correctness check assert F.allclose(r1, r5) assert F.allclose(r2, r6) assert F.allclose(r3, r7) assert F.allclose(r4, r8) assert F.allclose(grad_edata_hm, grad_edata_ht)