Add reduce option to decov

chainer/functions/loss/decov.py

@@ -10,9 +10,14 @@ class DeCov(function.Function):
 
     """DeCov loss (https://arxiv.org/abs/1511.06068)"""
 
-    def __init__(self):
+    def __init__(self, reduce='half_squared_sum'):
         self.h_centered = None
         self.covariance = None
+        if reduce not in ('half_squared_sum', 'no'):
+            raise ValueError(
+                "only 'half_squared_sum' and 'no' are valid "
+                "for 'reduce', but '%s' is given" % reduce)
+        self.reduce = reduce
 
     def check_type_forward(self, in_types):
         type_check.expect(in_types.size() == 1)
@@ -31,32 +36,56 @@ def forward(self, inputs):
         self.covariance = self.h_centered.T.dot(self.h_centered)
         xp.fill_diagonal(self.covariance, 0.0)
         self.covariance /= len(h)
-        cost = xp.vdot(self.covariance, self.covariance) * h.dtype.type(0.5)
-        return utils.force_array(cost),
+        if self.reduce == 'half_squared_sum':
+            cost = xp.vdot(self.covariance, self.covariance)
+            cost *= h.dtype.type(0.5)
+            return utils.force_array(cost),
+        else:
+            return self.covariance,
 
     def backward(self, inputs, grad_outputs):
+        xp = cuda.get_array_module(*inputs)
         h, = inputs
         gcost, = grad_outputs
         gcost_div_n = gcost / gcost.dtype.type(len(h))
-
-        gh = 2.0 * self.h_centered.dot(self.covariance)
-        gh *= gcost_div_n
+        if self.reduce == 'half_squared_sum':
+            gh = 2.0 * self.h_centered.dot(self.covariance)
+            gh *= gcost_div_n
+        else:
+            xp.fill_diagonal(gcost_div_n, 0.0)
+            gh = self.h_centered.dot(gcost_div_n + gcost_div_n.T)
         return gh,
 
 
-def decov(h):
+def decov(h, reduce='half_squared_sum'):
     """Computes the DeCov loss of ``h``
 
+    The output is a variable whose value depends on the value of
+    the option ``reduce``. If it is ``'no'``, it holds a matrix
+    whose size is same as the number of columns of ``y``.
+    If it is ``'half_squared_sum'``, it holds the half of the
+    squared Frobenius norm (i.e. squared of the L2 norm of a matrix flattened
+    to a vector) of the matrix.
+
     Args:
         h (Variable): Variable holding a matrix where the first dimension
             corresponds to the batches.
+        recude (str): Reduction option. Its value must be either
+            ``'half_squared_sum'`` or ``'no'``.
+            Otherwise, :class:`ValueError` is raised.
 
     Returns:
-        Variable: A variable holding a scalar of the DeCov loss.
+        ~chainer.Variable:
+            A variable holding a scalar of the DeCov loss.
+            If ``reduce`` is ``'no'``, the output variable holds
+            2-dimensional array matrix of shape ``(N, N)`` where
+            ``N`` is the number of columns of ``y``.
+            If it is ``'half_squared_sum'``, the output variable
+            holds a scalar value.
 
     .. note::
 
        See https://arxiv.org/abs/1511.06068 for details.
 
     """
-    return DeCov()(h)
+    return DeCov(reduce)(h)

tests/chainer_tests/functions_tests/loss_tests/test_decov.py

@@ -12,36 +12,51 @@
 from chainer.testing import condition
 
 
+def _deconv(h):
+    h = cuda.to_cpu(h)
+    h_mean = h.mean(axis=0)
+    N, M = h.shape
+    loss_expect = numpy.zeros((M, M), dtype=numpy.float32)
+    for i in six.moves.range(M):
+        for j in six.moves.range(M):
+            if i != j:
+                for n in six.moves.range(N):
+                    loss_expect[i, j] += (h[n, i] - h_mean[i]) * (
+                        h[n, j] - h_mean[j])
+    return loss_expect / N
+
+
+@testing.parameterize(
+    {'reduce': 'half_squared_sum'},
+    {'reduce': 'no'}
+)
 class TestDeCov(unittest.TestCase):
 
     def setUp(self):
         self.h = numpy.random.uniform(-1, 1, (4, 3)).astype(numpy.float32)
+        if self.reduce == 'half_squared_sum':
+            gloss_shape = ()
+        else:
+            gloss_shape = (3, 3)
+        self.gloss = numpy.random.uniform(
+            -1, 1, gloss_shape).astype(numpy.float32)
 
     def check_forward(self, h_data):
         h = chainer.Variable(h_data)
-        loss = functions.decov(h)
-        self.assertEqual(loss.data.shape, ())
+        loss = functions.decov(h, self.reduce)
+        self.assertEqual(loss.shape, self.gloss.shape)
         self.assertEqual(loss.data.dtype, numpy.float32)
-        loss_value = float(loss.data)
+        loss_value = cuda.to_cpu(loss.data)
 
         # Compute expected value
         h_data = cuda.to_cpu(h_data)
-        h_mean = h_data.mean(axis=0)
-        N = h_data.shape[0]
-
-        loss_expect = 0
-        for i in six.moves.range(h_data.shape[1]):
-            for j in six.moves.range(h_data.shape[1]):
-                ij_loss = 0.
-                if i != j:
-                    for n in six.moves.range(N):
-                        ij_loss += (h_data[n, i] - h_mean[i]) * (
-                            h_data[n, j] - h_mean[j])
-                    ij_loss /= N
-                loss_expect += ij_loss ** 2
-        loss_expect *= 0.5
-
-        self.assertAlmostEqual(loss_expect, loss_value, places=5)
+
+        loss_expect = _deconv(h_data)
+        if self.reduce == 'half_squared_sum':
+            loss_expect = (loss_expect ** 2).sum() * 0.5
+
+        numpy.testing.assert_allclose(
+            loss_expect, loss_value, rtol=1e-4, atol=1e-4)
 
     @condition.retry(3)
     def test_forward_cpu(self):
@@ -52,31 +67,54 @@ def test_forward_cpu(self):
     def test_forward_gpu(self):
         self.check_forward(cuda.to_gpu(self.h))
 
-    def check_backward(self, h_data):
+    def check_backward(self, h_data, gloss_data):
         gradient_check.check_backward(
-            functions.DeCov(), (h_data,), None, eps=0.02, atol=1e-3)
+            functions.DeCov(self.reduce), (h_data,), gloss_data,
+            eps=0.02, atol=1e-3)
 
-    def check_type(self, h_data):
+    def check_type(self, h_data, gloss_data):
         h = chainer.Variable(h_data)
-        loss = functions.decov(h)
+        loss = functions.decov(h, self.reduce)
+        loss.grad = gloss_data
         loss.backward()
         self.assertEqual(h_data.dtype, h.grad.dtype)
 
     @condition.retry(3)
     def test_backward_cpu(self):
-        self.check_backward(self.h)
+        self.check_backward(self.h, self.gloss)
 
     def test_backward_type_cpu(self):
-        self.check_type(self.h)
+        self.check_type(self.h, self.gloss)
 
     @attr.gpu
     def test_backward_type_gpu(self):
-        self.check_type(cuda.to_gpu(self.h))
+        self.check_type(cuda.to_gpu(self.h),
+                        cuda.to_gpu(self.gloss))
 
     @attr.gpu
     @condition.retry(3)
     def test_backward_gpu(self):
-        self.check_backward(cuda.to_gpu(self.h))
+        self.check_backward(cuda.to_gpu(self.h),
+                            cuda.to_gpu(self.gloss))
+
+
+class TestDeconvInvalidReductionOption(unittest.TestCase):
+
+    def setUp(self):
+        self.h = numpy.random.uniform(-1, 1, (4, 3)).astype(numpy.float32)
+
+    def check_invalid_option(self, xp):
+        h = xp.asarray(self.h)
+
+        with self.assertRaises(ValueError):
+            functions.decov(h, 'invalid_option')
+
+    def test_invalid_option_cpu(self):
+        self.check_invalid_option(numpy)
+
+    @attr.gpu
+    def test_invalid_option_gpu(self):
+        self.check_invalid_option(cuda.cupy)
 
 
 testing.run_module(__name__, __file__)