Repository: chainer/chainer

Issue:
Grad is wrong when using `WeightDecay` of Optimizer Hook and loss scaling
Grad is wrong when using `WeightDecay` and loss scaling.
The cause is that `WeightDecay` is performed before dividing by loss scale.

`WeightDecay` of optimizer hook is performed here:
https://github.com/chainer/chainer/blob/master/chainer/optimizer.py#L810

Dividing by loss scale is here:
https://github.com/chainer/chainer/blob/master/chainer/optimizer.py#L220

Chainer version: v6.0.0 and v7.0.0a1
Grad is wrong when using `WeightDecay` of Optimizer Hook and loss scaling
Grad is wrong when using `WeightDecay` and loss scaling.
The cause is that `WeightDecay` is performed before dividing by loss scale.

`WeightDecay` of optimizer hook is performed here:
https://github.com/chainer/chainer/blob/master/chainer/optimizer.py#L810

Dividing by loss scale is here:
https://github.com/chainer/chainer/blob/master/chainer/optimizer.py#L220

Chainer version: v6.0.0 and v7.0.0a1


Discussion:

Resolving patch:
diff --git a/chainer/optimizer_hooks/weight_decay.py b/chainer/optimizer_hooks/weight_decay.py
--- a/chainer/optimizer_hooks/weight_decay.py
+++ b/chainer/optimizer_hooks/weight_decay.py
@@ -42,9 +42,12 @@ def __call__(self, rule, param):
         if p is None or g is None:
             return
         with chainer.using_device(param.device):
+            rate = self.rate
+            if param._loss_scale is not None:
+                rate *= param._loss_scale
             if param.device.xp is cuda.cupy:
                 kernel = cuda.elementwise(
                     'T p, T decay', 'T g', 'g += decay * p', 'weight_decay')
-                kernel(p, self.rate, g)
+                kernel(p, rate, g)
             else:
-                g += self.rate * p
+                g += rate * p


Test patch:
diff --git a/tests/chainer_tests/optimizer_hooks_tests/test_weight_decay.py b/tests/chainer_tests/optimizer_hooks_tests/test_weight_decay.py
--- a/tests/chainer_tests/optimizer_hooks_tests/test_weight_decay.py
+++ b/tests/chainer_tests/optimizer_hooks_tests/test_weight_decay.py
@@ -3,6 +3,7 @@
 import numpy as np
 
 import chainer
+import chainer.functions as F
 import chainer.initializers as I
 from chainer import optimizer_hooks
 from chainer import optimizers
@@ -50,4 +51,36 @@ def test_weight_decay_gpu(self):
         self.check_weight_decay()
 
 
+@testing.inject_backend_tests(
+    None,
+    # CPU tests
+    [{}]
+    # GPU tests
+    + testing.product({
+        'use_cuda': [True],
+    })
+)
+class TestWeightDecayLossScale(unittest.TestCase):
+
+    def test_weight_decay_loss_scale(self, backend_config):
+        a = self._updated_array(backend_config, None)
+        b = self._updated_array(backend_config, loss_scale=4.)
+        testing.assert_allclose(a, b)
+
+    def _updated_array(self, backend_config, loss_scale):
+        arr = np.arange(3, dtype=np.float32)
+        param = chainer.Parameter(arr)
+        link = chainer.Link()
+        with link.init_scope():
+            link.p = param
+        link.to_device(backend_config.device)
+        opt = optimizers.SGD(lr=1)
+        opt.setup(link)
+        opt.add_hook(optimizer_hooks.WeightDecay(1/8.))
+        loss = F.sum(link.p ** 3)
+        loss.backward(loss_scale=loss_scale)
+        opt.update()
+        return link.p.array
+
+
 testing.run_module(__name__, __file__)