Repository: chainer/chainer

Issue:
Implement GPU mode for hierarchical softmax



Discussion:

Resolving patch:
diff --git a/chainer/functions/hierarchical_softmax.py b/chainer/functions/hierarchical_softmax.py
--- a/chainer/functions/hierarchical_softmax.py
+++ b/chainer/functions/hierarchical_softmax.py
@@ -1,6 +1,7 @@
 import numpy
 import six
 
+from chainer import cuda
 from chainer import function
 from chainer.utils import type_check
 
@@ -163,6 +164,112 @@ def _backward_cpu_one(self, x, t, gloss):
         self.gW[path] += gw
         return gx
 
+    def to_gpu(self, device=None):
+        function.Function.to_gpu(self, device)
+
+        n_vocab = max(self.paths.keys()) + 1
+        paths = cuda.to_gpu(numpy.concatenate(
+            [self.paths[i] for i in range(n_vocab) if i in self.paths]))
+        codes = cuda.to_gpu(numpy.concatenate(
+            [self.codes[i] for i in range(n_vocab) if i in self.codes]))
+
+        begins = numpy.empty((n_vocab + 1,), dtype=numpy.int32)
+        begins[0] = 0
+        for i in range(0, n_vocab):
+            length = len(self.paths[i]) if i in self.paths else 0
+            begins[i + 1] = begins[i] + length
+
+        self.paths = paths
+        self.codes = codes
+        self.begins = cuda.to_gpu(begins)
+
+    def forward_gpu(self, inputs):
+        x, t = inputs
+
+        max_length = cuda.reduce(
+            'int* t, int* begins', 'begins[t[i] + 1] - begins[t[i]]',
+            'max(a,b)', '0', 'binary_hierarchical_softmax_max_length',
+            numpy.int32
+        )(t, self.begins)
+        max_length = cuda.to_cpu(max_length)[()]
+
+        length = max_length * x.shape[0]
+        ls = cuda.empty((length,), dtype=numpy.float32)
+        n_in = x.shape[1]
+        wxy = cuda.empty((length,), dtype=numpy.float32)
+        cuda.elementwise(
+            '''float* ls, float* wxy, const float* x, const float* w,
+            const int* ts, const int* paths, const float* codes,
+            const int* begins, int c, int max_length''',
+            '''
+            int ind = i / max_length;
+            int offset = i - ind * max_length;
+            int t = ts[ind];
+
+            int begin = begins[t];
+            int length = begins[t + 1] - begins[t];
+
+            if (offset < length) {
+              int p = begin + offset;
+              int node = paths[p];
+
+              x = &x[ind * c];
+
+              float wx = 0;
+              for (int j = 0; j < c; ++j) {
+                wx += w[node * c + j] * x[j];
+              }
+              wxy[i] = wx * codes[p];
+              ls[i] = log(1 + exp(-wxy[i]));
+            } else {
+              ls[i] = 0;
+            }
+            ''',
+            'binary_hierarchical_softmax_forward'
+        )(ls, wxy, x, self.W, t, self.paths, self.codes, self.begins,
+          n_in, max_length)
+        self.max_length = max_length
+        self.wxy = wxy
+        return cuda.gpuarray.sum(ls),
+
+    def backward_gpu(self, inputs, loss):
+        x, t = inputs
+        gloss, = loss
+
+        n_in = x.shape[1]
+        gx = cuda.zeros_like(x)
+        cuda.elementwise(
+            '''const float* wxy, float* gx, float* gw, const float* x,
+            const float* w, const int* ts, const int* paths,
+            const float* codes, const int* begins,
+            const float* gloss, int c, int max_length''',
+            '''
+            int ind = i / max_length;
+            int offset = i - ind * max_length;
+            int t = ts[ind];
+
+            int begin = begins[t];
+            int length = begins[t + 1] - begins[t];
+
+            if (offset < length) {
+              int p = begin + offset;
+              int node = paths[p];
+              float code = codes[p];
+              gx = &gx[ind * c];
+              x = &x[ind * c];
+
+              float g = -*gloss * code / (1.0 + exp(wxy[i]));
+              for (int j = 0; j < c; ++j) {
+                atomicAdd(gx + j, g * w[node * c + j]);
+                atomicAdd(gw + node * c + j, g * x[j]);
+              }
+            }
+            ''',
+            'binary_hierarchical_softmax_bwd'
+        )(self.wxy, gx, self.gW, x, self.W, t, self.paths, self.codes,
+          self.begins, gloss, n_in, self.max_length)
+        return gx, None
+
 
 def create_huffman_tree(word_counts):
     """Make a huffman tree from a dictionary containing word counts.


Test patch:
diff --git a/tests/functions_tests/test_hierarchical_softmax.py b/tests/functions_tests/test_hierarchical_softmax.py
--- a/tests/functions_tests/test_hierarchical_softmax.py
+++ b/tests/functions_tests/test_hierarchical_softmax.py
@@ -7,9 +7,14 @@
 from chainer import functions
 from chainer import gradient_check
 from chainer import testing
+from chainer.testing import attr
 from chainer.testing import condition
 
 
+if cuda.available:
+    cuda.init()
+
+
 class TestHuffmanTree(unittest.TestCase):
 
     def test_empty(self):
@@ -34,19 +39,41 @@ def setUp(self):
 
         self.W = self.func.W.copy()
 
-    @condition.retry(3)
-    def test_sum(self):
-        x = numpy.array([[1.0, 2.0, 3.0]], numpy.float32)
+    def check_sum(self, x, gpu=False):
         total = 0
         for i in range(5):
             t = numpy.array([i])
-            loss, = self.func.forward_cpu((x, t))
+            if gpu:
+                t = cuda.to_gpu(t)
+            loss, = self.func.forward((x, t))
             self.assertEqual(loss.dtype, numpy.float32)
             self.assertEqual(loss.shape, ())
-            total += numpy.exp(-loss)
+            total += numpy.exp(-cuda.to_cpu(loss))
         self.assertAlmostEqual(1.0, float(total), delta=1.0e-5)
 
-    def check_backward(self, x_data, t_data, y_grad, use_cudnn=True):
+    @condition.retry(3)
+    def test_sum_cpu(self):
+        x = numpy.array([[1.0, 2.0, 3.0]], numpy.float32)
+        self.check_sum(x)
+
+    @attr.gpu
+    @condition.retry(3)
+    def test_sum_gpu(self):
+        x = numpy.array([[1.0, 2.0, 3.0]], numpy.float32)
+        self.func.to_gpu()
+        self.check_sum(cuda.to_gpu(x), gpu=True)
+
+    @attr.gpu
+    def test_forward(self):
+        # TODO(unno): We need to test return values of forward function.
+        cpu_loss, = self.func.forward((self.x, self.t))
+        self.func.to_gpu()
+        gpu_loss, = self.func.forward((cuda.to_gpu(self.x),
+                                       cuda.to_gpu(self.t)))
+        gradient_check.assert_allclose(
+            cpu_loss, cuda.to_cpu(gpu_loss))
+
+    def check_backward(self, x_data, t_data, y_grad):
         x = chainer.Variable(x_data)
         t = chainer.Variable(t_data)
         y = self.func(x, t)
@@ -65,5 +92,13 @@ def check_backward(self, x_data, t_data, y_grad, use_cudnn=True):
     def test_backward_cpu(self):
         self.check_backward(self.x, self.t, self.gy)
 
+    @attr.gpu
+    @condition.retry(3)
+    def test_backward_gpu(self):
+        self.func.to_gpu()
+        self.check_backward(cuda.to_gpu(self.x),
+                            cuda.to_gpu(self.t),
+                            cuda.to_gpu(self.gy))
+
 
 testing.run_module(__name__, __file__)