Merge bca94b5 into e6a1f8a

chainer/functions/activation/leaky_relu.py

@@ -1,15 +1,21 @@
 from chainer import cuda
-from chainer import function
+from chainer import function_node
 from chainer.utils import type_check
 
 
-def _kern():
-    return cuda.elementwise(
-        'T cond, T x, T slope', 'T y',
-        'y = cond >= 0 ? x : (T)(slope * x)', 'lrelu')
+_kern = None
 
 
-class LeakyReLU(function.Function):
+def _get_kern():
+    global _kern
+    if _kern is None:
+        _kern = cuda.elementwise(
+            'T cond, T x, T slope', 'T y',
+            'y = cond >= 0 ? x : (T)(slope * x)', 'lrelu')
+    return _kern
+
+
+class LeakyReLU(function_node.FunctionNode):
 
     """Leaky rectifier unit."""
 
@@ -21,37 +27,61 @@ def check_type_forward(self, in_types):
         x_type, = in_types
         type_check.expect(x_type.dtype.kind == 'f')
 
-    def forward_cpu(self, x):
-        y = x[0].copy()
-        y[x[0] < 0] *= self.slope
+    def forward_cpu(self, inputs):
+        x, = inputs
+        y = x.copy()
+        y[x < 0] *= self.slope
         if self.slope >= 0:
-            self.retain_inputs(())
             self.retain_outputs((0,))
+        else:
+            self.retain_inputs((0,))
         return y,
 
-    def forward_gpu(self, x):
-        y = _kern()(x[0], x[0], self.slope)
+    def forward_gpu(self, inputs):
+        x, = inputs
+        y = _get_kern()(x, x, self.slope)
         if self.slope >= 0:
-            self.retain_inputs(())
             self.retain_outputs((0,))
+        else:
+            self.retain_inputs((0,))
         return y,
 
-    def backward_cpu(self, x, gy):
-        gx = gy[0].copy()
+    def backward(self, indexes, grad_outputs):
+        if self.slope >= 0:
+            x = None
+            y = self.get_retained_outputs()[0].data
+        else:
+            x = self.get_retained_inputs()[0].data
+            y = None
+        return _LeakyReLUGrad(x, y, self.slope).apply(grad_outputs)
+
+
+class _LeakyReLUGrad(function_node.FunctionNode):
+
+    def __init__(self, x, y, slope):
+        self.slope = slope
+        self.x = x
+        self.y = y
+
+    def forward_cpu(self, inputs):
+        gy, = inputs
+        gy = gy.copy()
         if self.slope >= 0:
-            y = self.output_data
-            gx[y[0] < 0] *= self.slope
+            gy[self.y < 0] *= self.slope
         else:
-            gx[x[0] < 0] *= self.slope
-        return gx,
+            gy[self.x < 0] *= self.slope
+        return gy,
 
-    def backward_gpu(self, x, gy):
+    def forward_gpu(self, inputs):
+        gy, = inputs
         if self.slope >= 0:
-            y = self.output_data
-            gx = _kern()(y[0], gy[0], self.slope)
+            gy = _get_kern()(self.y, gy, self.slope)
         else:
-            gx = _kern()(x[0], gy[0], self.slope)
-        return gx,
+            gy = _get_kern()(self.x, gy, self.slope)
+        return gy,
+
+    def backward(self, indexes, grad_outputs):
+        return _LeakyReLUGrad(self.x, self.y, self.slope).apply(grad_outputs)
 
 
 def leaky_relu(x, slope=0.2):
@@ -86,4 +116,4 @@ def leaky_relu(x, slope=0.2):
                [-0.40000001,  1.        ]], dtype=float32)
 
     """
-    return LeakyReLU(slope)(x)
+    return LeakyReLU(slope).apply((x,))[0]

tests/chainer_tests/functions_tests/activation_tests/test_leaky_relu.py

@@ -25,13 +25,14 @@ def setUp(self):
             if -0.05 < self.x[i] < 0.05:
                 self.x[i] = 0.5
         self.gy = numpy.random.uniform(-1, 1, self.shape).astype(self.dtype)
+        self.ggx = numpy.random.uniform(-1, 1, self.shape).astype(self.dtype)
         self.slope = random.random()
         self.check_forward_options = {}
         self.check_backward_options = {'dtype': numpy.float64}
         if self.dtype == numpy.float16:
             self.check_forward_options = {'atol': 1e-4, 'rtol': 1e-3}
             self.check_backward_options = {
-                'dtype': numpy.float64, 'atol': 5e-4, 'rtol': 5e-3}
+                'dtype': numpy.float64, 'atol': 5e-3, 'rtol': 5e-2}
 
     def check_forward(self, x_data):
         x = chainer.Variable(x_data)
@@ -56,8 +57,11 @@ def test_forward_gpu(self):
         self.check_forward(cuda.to_gpu(self.x))
 
     def check_backward(self, x_data, y_grad):
+        def f(x):
+            return functions.leaky_relu(x, self.slope)
+
         gradient_check.check_backward(
-            functions.LeakyReLU(self.slope), x_data, y_grad,
+            f, x_data, y_grad,
             **self.check_backward_options)
 
     @condition.retry(10)
@@ -69,5 +73,26 @@ def test_backward_cpu(self):
     def test_backward_gpu(self):
         self.check_backward(cuda.to_gpu(self.x), cuda.to_gpu(self.gy))
 
+    def check_double_backward(self, x_data, y_grad, x_grad_grad):
+        def f(x):
+            y = functions.leaky_relu(x, self.slope)
+            return y * y
+
+        gradient_check.check_double_backward(
+            f, x_data, y_grad, x_grad_grad,
+            **self.check_backward_options)
+
+    @condition.retry(10)
+    def test_double_backward_cpu(self):
+        self.check_double_backward(self.x, self.gy, self.ggx)
+
+    @attr.gpu
+    @condition.retry(10)
+    def test_double_backward_gpu(self):
+        self.check_double_backward(
+            cuda.to_gpu(self.x),
+            cuda.to_gpu(self.gy),
+            cuda.to_gpu(self.ggx))
+
 
 testing.run_module(__name__, __file__)