Repository: chainer/chainer

Issue:
cleargrad in optimizer hooks
If `cleargrad` method is used in optimizer hooks, the optimizer raises `TypeError` exception.
```
class Freeze(object):
    name = 'Freeze'

    def __call__(self, opt):
        for param in opt.target.params():
            param.cleargrad()


model = chainer.links.Linear(1, 1)
optimizer = chainer.optimizers.SGD()
optimizer.setup(model)
optimizer.add_hook(Freeze())

optimizer.update()  # TypeError: unsupported operand type(s) for *: 'float' and 'NoneType'
```
If we use `param.grad[:] = 0` instead of `param.cleargrad`, it works well.

This is because `GradientMethod`  replace 'None' only before calling hooks.
https://github.com/pfnet/chainer/blob/master/chainer/optimizer.py#L400-L404


Discussion:

Resolving patch:
diff --git a/chainer/optimizer.py b/chainer/optimizer.py
--- a/chainer/optimizer.py
+++ b/chainer/optimizer.py
@@ -369,6 +369,27 @@ class GradientMethod(Optimizer):
 
     """
 
+    def reallocate_cleared_grads(self):
+        """Reallocate gradients cleared by :meth:`~chainer.Variable.cleargrad`.
+
+        This method allocates arrays for all gradients which have :obj:`None`.
+        This method is called before and after every optimizer hook.
+        If an inheriting optimizer does not require this allocation,
+        the optimizer can override this method with a blank function.
+
+        """
+        for name, param in self.target.namedparams():
+            if param.grad is None:
+                with cuda.get_device(param.data):
+                    xp = cuda.get_array_module(param.data)
+                    param.grad = xp.zeros_like(param.data)
+
+    def call_hooks(self):
+        """Invokes hook functions in registration order."""
+        for hook in six.itervalues(self._hooks):
+            hook(self)
+            self.reallocate_cleared_grads()
+
     def update(self, lossfun=None, *args, **kwds):
         """Updates parameters based on a loss function or computed gradients.
 
@@ -395,13 +416,7 @@ def update(self, lossfun=None, *args, **kwds):
             loss.backward()
             del loss
 
-        # TODO(unno): Some optimizers can skip this process if they does not
-        # affect to a parameter when its gradient is zero.
-        for name, param in self.target.namedparams():
-            if param.grad is None:
-                with cuda.get_device(param.data):
-                    xp = cuda.get_array_module(param.data)
-                    param.grad = xp.zeros_like(param.data)
+        self.reallocate_cleared_grads()
 
         self.call_hooks()
         self.prepare()


Test patch:
diff --git a/tests/chainer_tests/test_optimizer.py b/tests/chainer_tests/test_optimizer.py
--- a/tests/chainer_tests/test_optimizer.py
+++ b/tests/chainer_tests/test_optimizer.py
@@ -380,6 +380,30 @@ def test_clip_grads_gpu(self):
         self.check_clip_grads()
 
 
+class TestCleargradHook(unittest.TestCase):
+
+    def setUp(self):
+        self.target = SimpleLink(
+            np.arange(6, dtype=np.float32).reshape(2, 3),
+            np.arange(3, -3, -1, dtype=np.float32).reshape(2, 3))
+
+    def check_cleargrad(self):
+        opt = optimizers.SGD(lr=1)
+        opt.setup(self.target)
+        opt.add_hook(CleargradHook(self))
+        opt.add_hook(DummyHook(self))
+
+        opt.update()
+
+    def test_cleargrad_cpu(self):
+        self.check_cleargrad()
+
+    @attr.gpu
+    def test_cleargrad_gpu(self):
+        self.target.to_gpu()
+        self.check_cleargrad()
+
+
 class DummyOptimizer(chainer.GradientMethod):
 
     def __init__(self, test):
@@ -403,6 +427,19 @@ def __call__(self, opt):
             self.test.assertIsNotNone(param.grad)
 
 
+class CleargradHook(object):
+
+    name = 'Cleargrad'
+
+    def __init__(self, _):
+        pass
+
+    def __call__(self, opt):
+        for param in opt.target.params():
+            # Clear all grads
+            param.cleargrad()
+
+
 class TestGradientMethodClearGrads(unittest.TestCase):
 
     def setUp(self):