mirror of
https://github.com/dmlc/dgl.git
synced 2026-06-04 19:44:23 +08:00
269 lines
9.6 KiB
Python
269 lines
9.6 KiB
Python
import backend as F
|
|
|
|
import dgl
|
|
import numpy as np
|
|
import pytest
|
|
import scipy.sparse as ssp
|
|
from utils import parametrize_idtype
|
|
|
|
if F.backend_name == "pytorch":
|
|
import torch
|
|
|
|
torch.backends.cuda.matmul.allow_tf32 = False
|
|
|
|
|
|
def _random_simple_graph(
|
|
idtype, dtype, ctx, M, N, max_nnz, srctype, dsttype, etype
|
|
):
|
|
src = np.random.randint(0, M, (max_nnz,))
|
|
dst = np.random.randint(0, N, (max_nnz,))
|
|
val = np.random.randn(max_nnz)
|
|
a = ssp.csr_matrix((val, (src, dst)), shape=(M, N))
|
|
a.sum_duplicates()
|
|
a = a.tocoo()
|
|
# shuffle edges
|
|
perm = np.random.permutation(a.nnz)
|
|
row = a.row[perm]
|
|
col = a.col[perm]
|
|
val = a.data[perm]
|
|
a = ssp.csr_matrix((val, (row, col)), shape=(M, N))
|
|
|
|
A = dgl.heterograph(
|
|
{
|
|
(srctype, etype, dsttype): (
|
|
F.copy_to(F.tensor(row, dtype=idtype), ctx),
|
|
F.copy_to(F.tensor(col, dtype=idtype), ctx),
|
|
)
|
|
},
|
|
num_nodes_dict={srctype: a.shape[0], dsttype: a.shape[1]},
|
|
)
|
|
A.edata["w"] = F.copy_to(F.tensor(val, dtype=dtype), ctx)
|
|
return a, A
|
|
|
|
|
|
@parametrize_idtype
|
|
@pytest.mark.parametrize("dtype", [F.float32, F.float64])
|
|
@pytest.mark.parametrize("return_edge_ids", [True, False])
|
|
def test_csrmm(idtype, dtype, return_edge_ids):
|
|
a, A = _random_simple_graph(
|
|
idtype, dtype, F.ctx(), 500, 600, 9000, "A", "B", "AB"
|
|
)
|
|
b, B = _random_simple_graph(
|
|
idtype, dtype, F.ctx(), 600, 700, 9000, "B", "C", "BC"
|
|
)
|
|
C, C_weights = dgl._sparse_ops._csrmm(
|
|
A._graph, A.edata["w"], B._graph, B.edata["w"], 2
|
|
)
|
|
C_adj = C.adjacency_matrix_scipy(0, False, "csr", return_edge_ids)
|
|
C_adj.data = F.asnumpy(C_weights)
|
|
C_adj = F.tensor(C_adj.todense(), dtype=dtype)
|
|
c = F.tensor((a * b).todense(), dtype=dtype)
|
|
assert F.allclose(C_adj, c)
|
|
|
|
|
|
@parametrize_idtype
|
|
@pytest.mark.parametrize("dtype", [F.float32, F.float64])
|
|
@pytest.mark.parametrize("num_vtypes", [1, 2])
|
|
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)
|
|
|
|
|
|
@parametrize_idtype
|
|
@pytest.mark.parametrize("dtype", [F.float32, F.float64])
|
|
@pytest.mark.parametrize("return_edge_ids", [True, False])
|
|
def test_csrsum(idtype, dtype, return_edge_ids):
|
|
a, A = _random_simple_graph(
|
|
idtype, dtype, F.ctx(), 500, 600, 9000, "A", "B", "AB"
|
|
)
|
|
b, B = _random_simple_graph(
|
|
idtype, dtype, F.ctx(), 500, 600, 9000, "A", "B", "AB"
|
|
)
|
|
C, C_weights = dgl._sparse_ops._csrsum(
|
|
[A._graph, B._graph], [A.edata["w"], B.edata["w"]]
|
|
)
|
|
C_adj = C.adjacency_matrix_scipy(0, False, "csr", return_edge_ids)
|
|
C_adj.data = F.asnumpy(C_weights)
|
|
C_adj = F.tensor(C_adj.todense(), dtype=dtype)
|
|
c = F.tensor((a + b).todense(), dtype=dtype)
|
|
assert F.allclose(C_adj, c)
|
|
|
|
|
|
@parametrize_idtype
|
|
@pytest.mark.parametrize("dtype", [F.float32, F.float64])
|
|
@pytest.mark.parametrize("nelems", [1, 2])
|
|
def test_csrsum_backward(idtype, dtype, nelems):
|
|
a, A = _random_simple_graph(idtype, dtype, F.ctx(), 3, 4, 6, "A", "B", "AB")
|
|
b, B = _random_simple_graph(idtype, dtype, F.ctx(), 3, 4, 6, "A", "B", "AB")
|
|
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():
|
|
if nelems == 2:
|
|
# Test for two element case
|
|
C = dgl.adj_sum_graph([A, B], "w")
|
|
assert C.canonical_etypes == A.canonical_etypes
|
|
C_dense = np.zeros((3, 4))
|
|
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 = 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
|
|
)
|
|
elif nelems == 1:
|
|
# Test for single element case
|
|
C = dgl.adj_sum_graph([A], "w")
|
|
assert C.canonical_etypes == A.canonical_etypes
|
|
C_dense = np.zeros((3, 4))
|
|
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 = a_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]
|
|
A_spspmm_grad = F.asnumpy(F.grad(A.edata["w"]))
|
|
assert np.allclose(
|
|
a_dense_grad, A_spspmm_grad, rtol=1e-4, atol=1e-4
|
|
)
|
|
|
|
|
|
@parametrize_idtype
|
|
@pytest.mark.parametrize("dtype", [F.float32, F.float64])
|
|
@pytest.mark.parametrize("A_nnz", [9000, 0])
|
|
@pytest.mark.parametrize("B_nnz", [9000, 0])
|
|
def test_csrmask(idtype, dtype, A_nnz, B_nnz):
|
|
a, A = _random_simple_graph(
|
|
idtype, dtype, F.ctx(), 500, 600, A_nnz, "A", "B", "AB"
|
|
)
|
|
b, B = _random_simple_graph(
|
|
idtype, dtype, F.ctx(), 500, 600, B_nnz, "A", "B", "AB"
|
|
)
|
|
C = dgl._sparse_ops._csrmask(A._graph, A.edata["w"], B._graph)
|
|
B_row, B_col = B.edges(order="eid")
|
|
B_row = F.asnumpy(B_row)
|
|
B_col = F.asnumpy(B_col)
|
|
c = F.tensor(a.todense()[B_row, B_col], dtype)
|
|
assert F.allclose(C, c)
|
|
|
|
|
|
@parametrize_idtype
|
|
@pytest.mark.parametrize("dtype", [F.float32, F.float64])
|
|
def test_csrmask_backward(idtype, dtype):
|
|
a, A = _random_simple_graph(idtype, dtype, F.ctx(), 3, 4, 6, "A", "B", "AB")
|
|
b, B = _random_simple_graph(idtype, dtype, F.ctx(), 3, 4, 6, "A", "B", "AB")
|
|
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))
|
|
|
|
A.edata["w"] = F.attach_grad(A.edata["w"])
|
|
|
|
with F.record_grad():
|
|
# Test for two element case
|
|
C1 = F.csrmask(A._graph, A.edata["w"], B._graph)
|
|
if dgl.backend.backend_name == "tensorflow":
|
|
import tensorflow as tf
|
|
|
|
C2 = tf.gather_nd(a_dense, tf.stack([B_row, B_col], 1))
|
|
else:
|
|
C2 = a_dense[B_row, B_col]
|
|
assert F.allclose(C1, C2, rtol=1e-4, atol=1e-4)
|
|
|
|
F.backward(F.reduce_sum(C1) + F.reduce_sum(C2))
|
|
a_dense_grad = F.asnumpy(F.grad(a_dense))[A_row, A_col]
|
|
A_spspmm_grad = F.asnumpy(F.grad(A.edata["w"]))
|
|
assert np.allclose(a_dense_grad, A_spspmm_grad, rtol=1e-4, atol=1e-4)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_csrmm(F.int32, F.float32)
|
|
test_csrmm(F.int64, F.float32)
|
|
test_csrsum(F.int32, F.float32)
|
|
test_csrsum(F.int64, F.float32)
|
|
test_csrmask(F.int32, F.float32, 9000, 9000)
|
|
test_csrmask(F.int64, F.float32, 9000, 0)
|
|
test_csrmask(F.int32, F.float32, 0, 9000)
|
|
test_csrmask(F.int64, F.float32, 0, 0)
|
|
test_csrmm_backward(F.int32, F.float32, 1)
|
|
test_csrmm_backward(F.int64, F.float32, 1)
|
|
test_csrmm_backward(F.int32, F.float32, 2)
|
|
test_csrmm_backward(F.int64, F.float32, 2)
|
|
test_csrsum_backward(F.int32, F.float32, 1)
|
|
test_csrsum_backward(F.int64, F.float32, 1)
|
|
test_csrsum_backward(F.int32, F.float32, 2)
|
|
test_csrsum_backward(F.int64, F.float32, 2)
|
|
test_csrmask_backward(F.int32, F.float32)
|
|
test_csrmask_backward(F.int64, F.float32)
|