Merge pull request #5871 from grafi-tt/bug/fix-ideep-relu

Do not retain input on iDeep `F.relu`

chainer/functions/activation/relu.py

@@ -1,7 +1,6 @@
 import numpy
 
 import chainer
-from chainer import backend
 from chainer.backends import cuda
 from chainer.backends import intel64
 from chainer import function_node
@@ -24,7 +23,6 @@ class ReLU(function_node.FunctionNode):
     """Rectified Linear Unit."""
 
     _use_cudnn = False
-    _use_ideep = False
 
     def check_type_forward(self, in_types):
         type_check._argname(in_types, ('x',))
@@ -37,20 +35,17 @@ def forward_chainerx(self, inputs):
     def forward_cpu(self, inputs):
         if (intel64.should_use_ideep('>=auto')
                 and intel64.inputs_all_ready(inputs)):
-            # iDeep implementation
-            self._use_ideep = True
             return self.forward_ideep(inputs)
 
         x, = inputs
+        y = numpy.maximum(x, 0, dtype=x.dtype)
         self.retain_outputs((0,))
-        return utils.force_array(numpy.maximum(x, 0, dtype=x.dtype)),
+        return utils.force_array(y),
 
     def forward_ideep(self, inputs):
         x, = inputs
-        self.retain_inputs((0,))
-        self.retain_outputs((0,))
-
         y = intel64.ideep.relu.Forward(intel64.ideep.array(x))
+        self.retain_outputs((0,))
         return y,
 
     def forward_gpu(self, inputs):
@@ -69,20 +64,14 @@ def forward_gpu(self, inputs):
     def backward(self, indexes, grad_outputs):
         gy, = grad_outputs
         y, = self.get_retained_outputs()
-        if self._use_ideep:
-            # iDeep implementation
-            x, = self.get_retained_inputs()
-            return ReLUGradIdeep(x, y).apply((gy,))
+
         if chainer.should_use_cudnn('==always') and self._use_cudnn:
             # cuDNN implementation
             x, = self.get_retained_inputs()
-            return ReLUGradCudnn(x, y).apply((gy,))
-        # Generic implementation
-        return ReLUGrad2(y).apply((gy,))
+            return ReLUGradCudnn(x.array, y.array).apply((gy,))
 
-
-def _heaviside(x):
-    return (x > 0).astype(x.dtype)
+        # Generic implementation
+        return ReLUGrad2(y.array).apply((gy,))
 
 
 class ReLUGrad2(function_node.FunctionNode):
@@ -99,23 +88,33 @@ class ReLUGrad2(function_node.FunctionNode):
 
     def __init__(self, b):
         super(ReLUGrad2, self).__init__()
-        self.b = b.data
+        self.b = b
 
     def forward_cpu(self, inputs):
-        b = backend.from_chainerx(self.b)  # Workaround for ChainerX
-        y = (b > 0) * inputs[0]
-        return utils.force_array(y, dtype=y.dtype),
+        if (intel64.should_use_ideep('>=auto')
+                and intel64.inputs_all_ready(inputs)):
+            return self.forward_ideep(inputs)
+
+        gy, = inputs
+        gx = gy * (self.b > 0)
+        return utils.force_array(gx, dtype=gy.dtype),
+
+    def forward_ideep(self, inputs):
+        gy, = inputs
+        gx = intel64.ideep.relu.Backward(
+            intel64.ideep.array(self.b),
+            intel64.ideep.array(gy))
+        return gx,
 
     def forward_gpu(self, inputs):
-        b = backend.from_chainerx(self.b)  # Workaround for ChainerX
-        gx = _relu_grad2_kernel(b, inputs[0])
+        gx = _relu_grad2_kernel(self.b, inputs[0])
         return gx,
 
-    def backward(self, indexes, gy):
-        return gy[0] * _heaviside(self.b),
+    def backward(self, indexes, grad_outputs):
+        return ReLUGrad2(self.b).apply(grad_outputs)
 
 
-class ReLUGrad3Base(function_node.FunctionNode):
+class ReLUGradCudnn(function_node.FunctionNode):
     """Computes the gradient of the ReLU function.
 
     This function takes 3 variables a, b, and c, and
@@ -128,32 +127,17 @@ class ReLUGrad3Base(function_node.FunctionNode):
     """
 
     def __init__(self, x, y):
-        super(ReLUGrad3Base, self).__init__()
-        self.x = x.data
-        self.y = y.data
-
-    def backward(self, indexes, grad_outputs):
-        gy, = grad_outputs
-        ggx = gy * _heaviside(self.y)
-        return ggx,
-
-
-class ReLUGradCudnn(ReLUGrad3Base):
+        super(ReLUGradCudnn, self).__init__()
+        self.x = x
+        self.y = y
 
     def forward(self, inputs):
         assert chainer.should_use_cudnn('==always')
         gy, = inputs
         return cudnn.activation_backward(self.x, self.y, gy, _mode),
 
-
-class ReLUGradIdeep(ReLUGrad3Base):
-
-    def forward(self, inputs):
-        gy, = inputs
-        ggx = intel64.ideep.relu.Backward(
-            intel64.ideep.array(self.x),
-            intel64.ideep.array(gy))
-        return ggx,
+    def backward(self, indexes, grad_outputs):
+        return ReLUGrad2(self.y).apply(grad_outputs)
 
 
 def relu(x):