Repository: chainer/chainer

Issue:
Hierarchical softmax doesn't have `to_cpu`
Same problem as #276 



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
@@ -104,8 +104,20 @@ class BinaryHierarchicalSoftmax(function.Function):
     def __init__(self, in_size, tree):
         parser = TreeParser()
         parser.parse(tree)
-        self.paths = parser.get_paths()
-        self.codes = parser.get_codes()
+        paths = parser.get_paths()
+        codes = parser.get_codes()
+        n_vocab = max(paths.keys()) + 1
+
+        self.paths = numpy.concatenate(
+            [paths[i] for i in range(n_vocab) if i in paths])
+        self.codes = numpy.concatenate(
+            [codes[i] for i in range(n_vocab) if i in codes])
+        begins = numpy.empty((n_vocab + 1,), dtype=numpy.int32)
+        begins[0] = 0
+        for i in range(0, n_vocab):
+            length = len(paths[i]) if i in paths else 0
+            begins[i + 1] = begins[i] + length
+        self.begins = begins
 
         self.W = numpy.random.uniform(
             -1, 1, (parser.size(), in_size)).astype(numpy.float32)
@@ -139,10 +151,11 @@ def forward_cpu(self, args):
         return numpy.array(loss),
 
     def _forward_cpu_one(self, x, t):
-        assert t in self.paths
+        begin = self.begins[t]
+        end = self.begins[t + 1]
 
-        w = self.W[self.paths[t]]
-        wxy = w.dot(x) * self.codes[t]
+        w = self.W[self.paths[begin:end]]
+        wxy = w.dot(x) * self.codes[begin:end]
         loss = numpy.logaddexp(0.0, -wxy)  # == log(1 + exp(-wxy))
         return numpy.sum(loss)
 
@@ -155,10 +168,13 @@ def backward_cpu(self, args, loss):
         return gx, None
 
     def _backward_cpu_one(self, x, t, gloss):
-        path = self.paths[t]
+        begin = self.begins[t]
+        end = self.begins[t + 1]
+
+        path = self.paths[begin:end]
         w = self.W[path]
-        wxy = w.dot(x) * self.codes[t]
-        g = -gloss * self.codes[t] / (1.0 + numpy.exp(wxy))
+        wxy = w.dot(x) * self.codes[begin:end]
+        g = -gloss * self.codes[begin:end] / (1.0 + numpy.exp(wxy))
         gx = g.dot(w)
         gw = g.reshape((g.shape[0], 1)).dot(x.reshape(1, x.shape[0]))
         self.gW[path] += gw
@@ -167,21 +183,16 @@ def _backward_cpu_one(self, x, t, gloss):
     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]))
+        self.paths = cuda.to_gpu(self.paths, device)
+        self.codes = cuda.to_gpu(self.codes, device)
+        self.begins = cuda.to_gpu(self.begins, device)
 
-        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
+    def to_cpu(self):
+        function.Function.to_cpu(self)
 
-        self.paths = paths
-        self.codes = codes
-        self.begins = cuda.to_gpu(begins)
+        self.paths = cuda.to_cpu(self.paths)
+        self.codes = cuda.to_cpu(self.codes)
+        self.begins = cuda.to_cpu(self.begins)
 
     def forward_gpu(self, inputs):
         x, t = inputs


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
@@ -1,3 +1,4 @@
+import copy
 import unittest
 
 import numpy
@@ -107,5 +108,15 @@ def test_backward_gpu(self):
                             cuda.to_gpu(self.t),
                             cuda.to_gpu(self.gy))
 
+    @attr.gpu
+    def test_to_cpu(self):
+        f = copy.deepcopy(self.func)
+        self.func.to_gpu()
+        self.func.to_cpu()
+
+        self.assertTrue((f.begins == self.func.begins).all())
+        self.assertTrue((f.paths == self.func.paths).all())
+        self.assertTrue((f.codes == self.func.codes).all())
+
 
 testing.run_module(__name__, __file__)