Add tests for SparseGGNN
@@ -5,9 +5,11 @@ import pytest
from chainer_chemistry.config import MAX_ATOMIC_NUM
from chainer_chemistry.models.ggnn import GGNN
+from chainer_chemistry.models.ggnn import SparseGGNN
from chainer_chemistry.utils.extend import extend_node, extend_adj
from chainer_chemistry.utils.permutation import permute_adj
from chainer_chemistry.utils.permutation import permute_node
+from chainer_chemistry.utils.sparse_utils import _convert_to_sparse
atom_size = 5
out_dim = 4
@@ -17,9 +19,16 @@ num_edge_type = 3
@pytest.fixture
def model():
+ numpy.random.seed(0)
return GGNN(out_dim=out_dim, num_edge_type=num_edge_type)
+@pytest.fixture
+def sparse_model():
+ numpy.random.seed(0)
+ return SparseGGNN(out_dim=out_dim, num_edge_type=num_edge_type)
+
+
@pytest.fixture
def data():
numpy.random.seed(0)
@@ -34,27 +43,38 @@ def data():
return atom_data, adj_data, y_grad
-def check_forward(model, atom_data, adj_data):
- y_actual = cuda.to_cpu(model(atom_data, adj_data).data)
+def check_forward(model, *args):
+ numpy.random.seed(0) # reset seed to initialize model params consistently
+ y_actual = cuda.to_cpu(model(*args).data)
assert y_actual.shape == (batch_size, out_dim)
+ return y_actual
-def test_forward_cpu(model, data):
+def test_forward_cpu(model, sparse_model, data):
atom_data, adj_data = data[0], data[1]
- check_forward(model, atom_data, adj_data)
+ y_dense = check_forward(model, atom_data, adj_data)
+ # test for sparse data
+ y_sparse = check_forward(sparse_model, atom_data,
+ *_convert_to_sparse(adj_data))
+ numpy.testing.assert_allclose(
+ y_dense, y_sparse, atol=1e-4, rtol=1e-4)
@pytest.mark.gpu
-def test_forward_gpu(model, data):
+def test_forward_gpu(model, sparse_model, data):
atom_data, adj_data = cuda.to_gpu(data[0]), cuda.to_gpu(data[1])
model.to_gpu()
check_forward(model, atom_data, adj_data)
+ sparse_model.to_gpu()
+ check_forward(sparse_model, atom_data, *_convert_to_sparse(adj_data))
def test_backward_cpu(model, data):
atom_data, adj_data, y_grad = data
gradient_check.check_backward(model, (atom_data, adj_data), y_grad,
atol=1e-3, rtol=1e-3)
+ # there is no backward test for sparse model, because there will be no
+ # gradient for input data.
@pytest.mark.gpu
