Make NonbiasWeightDecay work again

chainerrl/optimizers/nonbias_weight_decay.py

@@ -11,26 +11,32 @@
 
 class NonbiasWeightDecay(object):
 
-    """Optimizer hook function for weight decay regularization.
+    """Weight decay only for non-bias parameters.
 
+    This hook can be used just like chainer.optimizer_hooks.WeightDecay except
+    that this hook does not apply weight decay to bias parameters.
+
+    This hook assumes that all the bias parameters have the name of "b". Any
+    parameter whose name is "b" is considered as a bias and excluded from
+    weight decay.
     """
     name = 'NonbiasWeightDecay'
+    call_for_each_param = True
+    timing = 'pre'
 
     def __init__(self, rate):
         self.rate = rate
 
-    def __call__(self, opt):
-        if cuda.available:
-            kernel = cuda.elementwise(
-                'T p, T decay', 'T g', 'g += decay * p', 'weight_decay')
-
-        rate = self.rate
-        for name, param in opt.target.namedparams():
-            if name == 'b' or name.endswith('/b'):
-                continue
-            p, g = param.array, param.grad
-            with cuda.get_device(p) as dev:
-                if int(dev) == -1:
-                    g += rate * p
-                else:
-                    kernel(p, rate, g)
+    def __call__(self, rule, param):
+        if param.name == 'b':
+            return
+        p, g = param.data, param.grad
+        if p is None or g is None:
+            return
+        with cuda.get_device_from_array(p) as dev:
+            if int(dev) == -1:
+                g += self.rate * p
+            else:
+                kernel = cuda.elementwise(
+                    'T p, T decay', 'T g', 'g += decay * p', 'weight_decay')
+                kernel(p, self.rate, g)

tests/optimizer_tests/test_nonbias_weight_decay.py

@@ -0,0 +1,43 @@
+from __future__ import unicode_literals
+from __future__ import print_function
+from __future__ import division
+from __future__ import absolute_import
+from builtins import *  # NOQA
+from future import standard_library
+standard_library.install_aliases()  # NOQA
+
+import unittest
+
+import chainer
+import chainer.links as L
+from chainer import testing
+import numpy as np
+
+import chainerrl
+
+
+@testing.parameterize(*testing.product(
+    {
+        'lr': [1.0, 0.1],
+        'weight_decay_rate': [0.1, 0.05]
+    }
+))
+class TestNonbiasWeightDecay(unittest.TestCase):
+
+    def test(self):
+
+        model = chainer.Chain(
+            a=L.Linear(1, 2, initialW=3, initial_bias=3),
+            b=chainer.Chain(c=L.Linear(2, 3, initialW=4, initial_bias=4)),
+        )
+        optimizer = chainer.optimizers.SGD(self.lr)
+        optimizer.setup(model)
+        optimizer.add_hook(
+            chainerrl.optimizers.NonbiasWeightDecay(
+                rate=self.weight_decay_rate))
+        optimizer.update(lambda: chainer.Variable(np.asarray(0.0)))
+        decay_factor = 1 - self.lr * self.weight_decay_rate
+        np.testing.assert_allclose(model.a.W.array, 3 * decay_factor)
+        np.testing.assert_allclose(model.a.b.array, 3)
+        np.testing.assert_allclose(model.b.c.W.array, 4 * decay_factor)
+        np.testing.assert_allclose(model.b.c.b.array, 4)