Repository: chainer/chainer

Issue:
Add an option to delete forward result in softmax cross entropy
When the number of output labels is large especially in NLP, forward result of softmax cross entropy `y` consumes a lot of memory. I want an option to delete this value after forward computation.



Discussion:

Resolving patch:
diff --git a/chainer/functions/loss/softmax_cross_entropy.py b/chainer/functions/loss/softmax_cross_entropy.py
--- a/chainer/functions/loss/softmax_cross_entropy.py
+++ b/chainer/functions/loss/softmax_cross_entropy.py
@@ -51,9 +51,10 @@ class SoftmaxCrossEntropy(function.Function):
 
     ignore_label = -1
 
-    def __init__(self, use_cudnn=True, normalize=True):
+    def __init__(self, use_cudnn=True, normalize=True, cache_score=True):
         self.use_cudnn = use_cudnn
         self.normalize = normalize
+        self.cache_score = cache_score
 
     def check_type_forward(self, in_types):
         type_check.expect(in_types.size() == 2)
@@ -82,7 +83,8 @@ def forward_cpu(self, inputs):
             self._check_input_values(x, t)
 
         log_y = softmax_log(x, False)
-        self.y = numpy.exp(log_y)
+        if self.cache_score:
+            self.y = numpy.exp(log_y)
         log_yd = numpy.rollaxis(log_y, 1)
         log_yd = log_yd.reshape(len(log_yd), -1)
 
@@ -106,7 +108,8 @@ def forward_gpu(self, inputs):
             self._check_input_values(x, t)
 
         log_y = softmax_log(x, self.use_cudnn)
-        self.y = cupy.exp(log_y)
+        if self.cache_score:
+            self.y = cupy.exp(log_y)
         if getattr(self, 'normalize', True):
             coeff = cupy.maximum(1, (t != self.ignore_label).sum())
         else:
@@ -125,20 +128,25 @@ def backward_cpu(self, inputs, grad_outputs):
         x, t = inputs
         gloss = grad_outputs[0]
         n_unit = t.size // len(t)
-        if self.y.ndim == 2:
-            gx = self.y.copy()
+        if hasattr(self, 'y'):
+            y = self.y.copy()
+        else:
+            log_y = softmax_log(x, self.use_cudnn)
+            y = numpy.exp(log_y)
+        if y.ndim == 2:
+            gx = y
             gx[six.moves.xrange(len(t)), numpy.maximum(t, 0)] -= 1
             gx *= (t != self.ignore_label).reshape((len(t), 1))
         else:
             # in the case where y.ndim is higher than 2,
             # we think that a current implementation is inefficient
             # because it yields two provisional arrays for indexing.
-            gx = self.y.copy().reshape(self.y.shape[0], self.y.shape[1], -1)
+            gx = y.reshape(y.shape[0], y.shape[1], -1)
             fst_index = numpy.arange(t.size) // n_unit
             trd_index = numpy.arange(t.size) % n_unit
             gx[fst_index, numpy.maximum(t.ravel(), 0), trd_index] -= 1
             gx *= (t != self.ignore_label).reshape((len(t), 1, -1))
-            gx = gx.reshape(self.y.shape)
+            gx = gx.reshape(y.shape)
 
         gx *= gloss * self._coeff
         return gx, None
@@ -146,6 +154,11 @@ def backward_cpu(self, inputs, grad_outputs):
     def backward_gpu(self, inputs, grad_outputs):
         cupy = cuda.cupy
         x, t = inputs
+        if hasattr(self, 'y'):
+            y = self.y
+        else:
+            y = softmax_log(x, self.use_cudnn)
+            cupy.exp(y, out=y)
         gloss = grad_outputs[0]
         n_unit = t.size // len(t)
         coeff = gloss * self._coeff
@@ -157,11 +170,12 @@ def backward_gpu(self, inputs, grad_outputs):
                gx = (t == -1) ? 0 : (coeff[0] * (y - (c == t)));
             ''',
             'softmax_crossent_bwd')(
-                self.y, cupy.expand_dims(t, 1), coeff, x.shape[1], n_unit)
+                y, cupy.expand_dims(t, 1), coeff, x.shape[1], n_unit)
         return gx, None
 
 
-def softmax_cross_entropy(x, t, use_cudnn=True, normalize=True):
+def softmax_cross_entropy(
+        x, t, use_cudnn=True, normalize=True, cache_score=True):
     """Computes cross entropy loss for pre-softmax activations.
 
     Args:
@@ -178,6 +192,9 @@ def softmax_cross_entropy(x, t, use_cudnn=True, normalize=True):
             determines the normalization constant. If true, this function
             normalizes the cross entropy loss across all instances. If else,
             it only normalizes along a batch size.
+        cache_score (bool): When it is ``True``, the function stores result
+            of forward computation to use it on backward computation. It
+            reduces computational cost though consumes more memories.
 
     Returns:
         Variable: A variable holding a scalar array of the cross entropy loss.
@@ -187,4 +204,4 @@ def softmax_cross_entropy(x, t, use_cudnn=True, normalize=True):
        This function is differentiable only by ``x``.
 
     """
-    return SoftmaxCrossEntropy(use_cudnn, normalize)(x, t)
+    return SoftmaxCrossEntropy(use_cudnn, normalize, cache_score)(x, t)


Test patch:
diff --git a/tests/chainer_tests/functions_tests/loss_tests/test_softmax_cross_entropy.py b/tests/chainer_tests/functions_tests/loss_tests/test_softmax_cross_entropy.py
--- a/tests/chainer_tests/functions_tests/loss_tests/test_softmax_cross_entropy.py
+++ b/tests/chainer_tests/functions_tests/loss_tests/test_softmax_cross_entropy.py
@@ -1,4 +1,3 @@
-import math
 import unittest
 
 import mock
@@ -16,28 +15,36 @@
 
 class TestSoftmaxCrossEntropy(unittest.TestCase):
 
+    shape = (4, 3)
     backward_atol = 1e-4
+    cache_score = True
 
     def setUp(self):
-        self.x = numpy.random.uniform(-1, 1, (4, 3)).astype(numpy.float32)
-        self.t = numpy.random.randint(0, 3, (4,)).astype(numpy.int32)
+        self.x = numpy.random.uniform(-1, 1, self.shape).astype(numpy.float32)
+        out_shape = (self.shape[0],) + self.shape[2:]
+        self.t = numpy.random.randint(0, 3, out_shape).astype(numpy.int32)
 
     def check_forward(self, x_data, t_data, use_cudnn=True):
         x = chainer.Variable(x_data)
         t = chainer.Variable(t_data)
-        loss = functions.softmax_cross_entropy(x, t, use_cudnn)
+        loss = functions.softmax_cross_entropy(
+            x, t, use_cudnn=use_cudnn, cache_score=self.cache_score)
         self.assertEqual(loss.data.shape, ())
         self.assertEqual(loss.data.dtype, numpy.float32)
+        self.assertEqual(hasattr(loss.creator, 'y'), self.cache_score)
         loss_value = float(cuda.to_cpu(loss.data))
 
         # Compute expected value
         loss_expect = 0.0
         count = 0
-        for i in six.moves.range(self.x.shape[0]):
-            if self.t[i] == -1:
+        x = numpy.rollaxis(self.x, 1, self.x.ndim).reshape(
+            (self.t.size, self.x.shape[1]))
+        t = self.t.ravel()
+        for xi, ti in six.moves.zip(x, t):
+            if ti == -1:
                 continue
-            log_z = numpy.ufunc.reduce(numpy.logaddexp, self.x[i])
-            loss_expect -= (self.x[i] - log_z)[self.t[i]]
+            log_z = numpy.ufunc.reduce(numpy.logaddexp, xi)
+            loss_expect -= (xi - log_z)[ti]
             count += 1
 
         if count == 0:
@@ -63,7 +70,8 @@ def test_forward_gpu_no_cudnn(self):
 
     def check_backward(self, x_data, t_data, use_cudnn=True):
         gradient_check.check_backward(
-            functions.SoftmaxCrossEntropy(use_cudnn),
+            functions.SoftmaxCrossEntropy(
+                use_cudnn=use_cudnn, cache_score=self.cache_score),
             (x_data, t_data), None, eps=0.02, atol=self.backward_atol)
 
     @condition.retry(3)
@@ -81,6 +89,11 @@ def test_backward_gpu_no_cudnn(self):
         self.check_backward(cuda.to_gpu(self.x), cuda.to_gpu(self.t), False)
 
 
+class TestSoftmaxCrossEntropyRemoveForward(TestSoftmaxCrossEntropy):
+
+    cache_score = False
+
+
 class TestSoftmaxCrossEntropyUnstable(TestSoftmaxCrossEntropy):
 
     backward_atol = 1e-3
@@ -96,34 +109,6 @@ def setUp(self):
         self.x = numpy.random.uniform(-1, 1, (4, 3, 2)).astype(numpy.float32)
         self.t = numpy.random.randint(0, 3, (4, 2)).astype(numpy.int32)
 
-    def check_forward(self, x_data, t_data, use_cudnn=True):
-        x = chainer.Variable(x_data)
-        t = chainer.Variable(t_data)
-        loss = functions.softmax_cross_entropy(
-            x, t, use_cudnn, normalize=True)
-        self.assertEqual(loss.data.shape, ())
-        self.assertEqual(loss.data.dtype, numpy.float32)
-        loss_value = float(cuda.to_cpu(loss.data))
-
-        # Compute expected value
-        y = numpy.exp(self.x)
-        loss_expect = 0.0
-        count = 0
-        for i in six.moves.range(y.shape[0]):
-            for k in six.moves.range(y.shape[2]):
-                if self.t[i, k] == -1:
-                    continue
-                loss_expect -= math.log(
-                    y[i, self.t[i, k], k] / y[i, :, k].sum())
-                count += 1
-
-        if count == 0:
-            loss_expect = 0.0
-        else:
-            loss_expect /= count
-
-        self.assertAlmostEqual(loss_expect, loss_value, places=4)
-
 
 class TestReplicatedSoftmaxCrossEntropy2(TestSoftmaxCrossEntropy):
 
@@ -132,29 +117,6 @@ def setUp(self):
             -1, 1, (4, 3, 2, 5)).astype(numpy.float32)
         self.t = numpy.random.randint(0, 3, (4, 2, 5)).astype(numpy.int32)
 
-    def check_forward(self, x_data, t_data, use_cudnn=True):
-        x = chainer.Variable(x_data)
-        t = chainer.Variable(t_data)
-        loss = functions.softmax_cross_entropy(
-            x, t, use_cudnn, normalize=False)
-        self.assertEqual(loss.data.shape, ())
-        self.assertEqual(loss.data.dtype, numpy.float32)
-        loss_value = float(cuda.to_cpu(loss.data))
-
-        # Compute expected value
-        y = numpy.exp(self.x)
-        loss_expect = 0.0
-        for i in six.moves.range(y.shape[0]):
-            for k in six.moves.range(y.shape[2]):
-                for l in six.moves.range(y.shape[3]):
-                    if self.t[i, k, l] == -1:
-                        continue
-                    loss_expect -= math.log(
-                        y[i, self.t[i, k, l], k, l] / y[i, :, k, l].sum())
-        loss_expect /= y.shape[0]
-
-        self.assertAlmostEqual(loss_expect, loss_value, places=4)
-
 
 class TestSoftmaxCrossEntropyWithIgnoreLabel(TestSoftmaxCrossEntropy):