def test_csrmm_backward(idtype, dtype, num_vtypes):
a, A = _random_simple_graph(idtype, dtype, F.ctx(), 3, 4, 6, 'A', 'B', 'AB')
b, B = _random_simple_graph(idtype, dtype, F.ctx(), 4, 3, 6, 'B', 'A' if num_vtypes == 1 else 'C', 'BA')
A_row, A_col = A.edges(order='eid')
B_row, B_col = B.edges(order='eid')
A_row = F.asnumpy(A_row)
A_col = F.asnumpy(A_col)
B_row = F.asnumpy(B_row)
B_col = F.asnumpy(B_col)
a_dense = F.attach_grad(F.tensor(a.todense(), dtype=dtype))
b_dense = F.attach_grad(F.tensor(b.todense(), dtype=dtype))
A.edata['w'] = F.attach_grad(A.edata['w'])
B.edata['w'] = F.attach_grad(B.edata['w'])
with F.record_grad():
C = dgl.adj_product_graph(A, B, 'w')
assert len(C.ntypes) == num_vtypes
assert len(C.etypes) == 1
C_dense = np.zeros((3, 3))
C_row, C_col = C.edges(order='eid')
C_row = F.asnumpy(C_row)
C_col = F.asnumpy(C_col)
C_dense[C_row, C_col] = F.asnumpy(C.edata['w'])
c_dense = F.matmul(a_dense, b_dense)
assert np.allclose(C_dense, F.asnumpy(c_dense), rtol=1e-4, atol=1e-4)
F.backward(F.reduce_sum(C.edata['w']) + F.reduce_sum(c_dense))
a_dense_grad = F.asnumpy(F.grad(a_dense))[A_row, A_col]
b_dense_grad = F.asnumpy(F.grad(b_dense))[B_row, B_col]
A_spspmm_grad = F.asnumpy(F.grad(A.edata['w']))
B_spspmm_grad = F.asnumpy(F.grad(B.edata['w']))
assert np.allclose(a_dense_grad, A_spspmm_grad, rtol=1e-4, atol=1e-4)
assert np.allclose(b_dense_grad, B_spspmm_grad, rtol=1e-4, atol=1e-4)
def test_khop_adj():
N = 20
feat = F.randn((N, 5))
g = dgl.DGLGraph(nx.erdos_renyi_graph(N, 0.3))
for k in range(3):
adj = F.tensor(dgl.khop_adj(g, k))
# use original graph to do message passing for k times.
g.ndata['h'] = feat
for _ in range(k):
g.update_all(fn.copy_u('h', 'm'), fn.sum('m', 'h'))
h_0 = g.ndata.pop('h')
# use k-hop adj to do message passing for one time.
h_1 = F.matmul(adj, feat)
assert F.allclose(h_0, h_1, rtol=1e-3, atol=1e-3)
def _test(feat_scale):
in_feat = 16 * feat_scale
out_feat = 8 * feat_scale
print("in/out feat", in_feat, out_feat)
E_per_rel = F.copy_to(
F.tensor([
50, 100, 20, 284, 89, 10, 82, 9200, 10, 20, 30, 100, 128, 20,
284, 89, 10, 82, 92, 10, 20, 30, 100, 1280, 20, 284, 89, 1000,
82, 92, 10, 2000, 30, 100, 128, 20, 284, 89, 10, 82, 92, 10,
20, 30
]), F.cpu())
E_per_rel *= n_edge_scale
num_rel = len(E_per_rel)
print('num_rel', num_rel)
W_per_len = F.copy_to(
F.full((num_rel, ), in_feat, dtype=F.dtype(E_per_rel)), F.cpu())
H_arr = []
W_arr = []
Out_arr = []
Out_grad_arr = []
for eid in range(num_rel):
H_arr.append(F.randn((E_per_rel[eid], in_feat)))
W_arr.append(F.randn((in_feat, out_feat)))
Out_arr.append(F.zeros((E_per_rel[eid], out_feat)))
Out_grad_arr.append(F.ones((E_per_rel[eid], out_feat)))
H = F.cat([h for h in H_arr], 0)
W = F.cat([w for w in W_arr], 0)
W_3D = W.reshape(num_rel, in_feat, out_feat)
Out = F.cat([out for out in Out_arr], 0)
Out_grad = F.cat([o for o in Out_grad_arr], 0)
print('H.shape', H.shape)
print('W.shape', W.shape)
print('W_3D.shape', W_3D.shape)
print('Out.shape', Out.shape)
etype_arr = []
for eid in range(num_rel):
etype_arr.append(
F.full((E_per_rel[eid], ), eid, dtype=F.dtype(E_per_rel)))
etypes = F.cat([etype for etype in etype_arr], 0)
#################################################################
# low-mem version using PyTorch operator
#################################################################
# forward pass
out = []
for i in range(len(E_per_rel)):
Hi = H_arr[i]
Wi = W_arr[i]
out.append(F.matmul(Hi, Wi))
out_low_mem = F.cat(out, 0)
# backward pass
H_grad = []
W_grad = []
for i in range(len(E_per_rel)):
Hi = H_arr[i]
Wi = W_arr[i]
Out_gradi = Out_grad_arr[i]
H_grad.append(F.matmul(Out_gradi, Wi.transpose(0, 1)))
W_grad.append(F.matmul(Hi.transpose(0, 1), Out_gradi))
Hgrad_low_mem = F.cat(H_grad, 0)
Wgrad_low_mem = F.cat(W_grad, 0)
Wgrad_low_mem = Wgrad_low_mem.reshape(num_rel, in_feat, out_feat)
#################################################################
# gather_mm where H sorted according to etype
#################################################################
seglen_A = E_per_rel
F.attach_grad(H)
F.attach_grad(W_3D)
with F.record_grad():
out_gmm_sorted = dgl.ops.segment_mm(H, W_3D, seglen_A)
F.backward(F.reduce_sum(out_gmm_sorted))
Hgrad_gmm_sorted = H.grad
Wgrad_gmm_sorted = W_3D.grad
#################################################################
# gather_mm where H is not sorted (backward not supported yet)
#################################################################
F.attach_grad(H)
F.attach_grad(W_3D)
with F.record_grad():
out_gmm_unsorted = dgl.ops.gather_mm(H, W_3D, idx_rhs=etypes)
F.backward(F.reduce_sum(out_gmm_unsorted))
Hgrad_gmm_unsorted = H.grad
Wgrad_gmm_unsorted = W_3D.grad
# correctness check
assert F.allclose(out_low_mem, out_gmm_sorted, atol=1e-3, rtol=1e-3)
assert F.allclose(Hgrad_low_mem,
Hgrad_gmm_sorted,
atol=1e-3,
rtol=1e-3)
assert F.allclose(Wgrad_low_mem,
Wgrad_gmm_sorted,
atol=1e-3,
rtol=1e-3)
assert F.allclose(out_low_mem, out_gmm_unsorted, atol=1e-3, rtol=1e-3)
assert F.allclose(Hgrad_low_mem,
Hgrad_gmm_unsorted,
atol=1e-3,
rtol=1e-3)
assert F.allclose(Wgrad_low_mem,
Wgrad_gmm_unsorted,
atol=1e-3,
rtol=1e-3)
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)
