Repository: chainer/chainer

Issue:
Support loss_scale on chainerx backprop



Discussion:

Resolving patch:
diff --git a/chainer/function_node.py b/chainer/function_node.py
--- a/chainer/function_node.py
+++ b/chainer/function_node.py
@@ -1014,10 +1014,6 @@ def grad(outputs, inputs, grad_outputs=None, grad_inputs=None, set_grad=False,
     n_chx_inputs = sum([False if x is None else x._has_chainerx_array
                         for x in inputs])
     if n_chx_inputs == len(inputs):
-        if loss_scale is not None:
-            raise ValueError(
-                'loss_scale is not supported on chainerx.grad interface')
-
         # Need to access the arrays to invoke the chainer grad function
         if grad_outputs:
             grad_outputs_chx = [x._data[0] for x in grad_outputs]
@@ -1025,12 +1021,16 @@ def grad(outputs, inputs, grad_outputs=None, grad_inputs=None, set_grad=False,
             grad_outputs_chx = []
         outputs_chx = [x._data[0] for x in outputs]
         inputs_chx = [x._data[0] for x in inputs]
+        # pybind has issues when converting opt<int> -> opt<float>
+        if loss_scale is not None:
+            loss_scale = float(loss_scale)
         grads = chainerx.grad(outputs_chx, inputs_chx,
                               backprop_id=None,
                               enable_double_backprop=enable_double_backprop,
                               set_grad=set_grad,
                               retain_grad=retain_grad,
-                              grad_outputs=grad_outputs_chx)
+                              grad_outputs=grad_outputs_chx,
+                              loss_scale=loss_scale)
 
         if grad_inputs:
             grads = [g+gi._data[0] for g, gi in zip(grads, grad_inputs)]
diff --git a/chainer/optimizer.py b/chainer/optimizer.py
--- a/chainer/optimizer.py
+++ b/chainer/optimizer.py
@@ -283,7 +283,8 @@ def __update(self, param):
             self._init_states(param_)
 
             # Apply loss scaling
-            if loss_scale is not None:
+            if (loss_scale is not None
+                    and not isinstance(param_.array, chainerx.ndarray)):
                 param_.grad /= loss_scale
 
         # Call update_core
diff --git a/chainer/variable.py b/chainer/variable.py
--- a/chainer/variable.py
+++ b/chainer/variable.py
@@ -1433,13 +1433,14 @@ def backward(self, retain_grad=False, enable_double_backprop=False,
             if retain_grad:
                 raise RuntimeError(
                     'retain_grad is not supported for ChainerX array.')
-            if loss_scale is not None:
-                raise RuntimeError(
-                    'loss_scale is not supported for ChainerX array.')
             arr = self._data[0]
             assert isinstance(arr, chainerx.ndarray)
+            # pybind has issues when converting int -> opt<float>
+            if loss_scale:
+                loss_scale = float(loss_scale)
             chainerx.backward(
-                arr, enable_double_backprop=enable_double_backprop)
+                arr, enable_double_backprop=enable_double_backprop,
+                loss_scale=loss_scale)
             return
 
         # Initialize error by 1, if this is a loss variable


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
@@ -71,6 +71,11 @@ def test_weight_decay(self, backend_config0,
     + testing.product({
         'use_cuda': [True],
     })
+    + [
+        # ChainerX
+        {'use_chainerx': True, 'chainerx_device': 'native:0'},
+        {'use_chainerx': True, 'chainerx_device': 'cuda:0'}
+    ]
 )
 class TestWeightDecayLossScale(unittest.TestCase):
 
diff --git a/tests/chainer_tests/test_function_node.py b/tests/chainer_tests/test_function_node.py
--- a/tests/chainer_tests/test_function_node.py
+++ b/tests/chainer_tests/test_function_node.py
@@ -964,8 +964,6 @@ def check_grad(self):
             raise
 
     def test_grad(self, backend_config):
-        if self.loss_scale and backend_config.xp is chainerx:
-            pytest.skip('chainerx.grad does not support loss_scale')
         self.use_device(backend_config.device)
         self.check_grad()
 
@@ -991,8 +989,6 @@ def check_double_grad(self):
             raise
 
     def test_double_grad(self, backend_config):
-        if self.loss_scale and backend_config.xp is chainerx:
-            pytest.skip('chainerx.grad does not support loss_scale')
         self.use_device(backend_config.device)
         self.check_double_grad()
 
@@ -1007,9 +1003,6 @@ def test_double_grad(self, backend_config):
         {'use_ideep': 'always'},
         {'use_cuda': True, 'cuda_device': 0},
         {'use_cuda': True, 'cuda_device': 1},
-        {'use_chainerx': True, 'chainerx_device': 'native:0'},
-        {'use_chainerx': True, 'chainerx_device': 'cuda:0'},
-        {'use_chainerx': True, 'chainerx_device': 'cuda:1'},
     ]
 )
 class TestGradSimple(GradTestBase, unittest.TestCase):
@@ -1033,6 +1026,34 @@ def expected_double_grad(self):
         return [ggrad]
 
 
+@testing.parameterize(*testing.product({
+    'loss_scale': [None, 1, 1.5, 2.5, 10],
+}))
+@testing.backend.inject_backend_tests(
+    None,
+    [
+        {'use_chainerx': True, 'chainerx_device': 'native:0'},
+        {'use_chainerx': True, 'chainerx_device': 'cuda:0'},
+        {'use_chainerx': True, 'chainerx_device': 'cuda:1'},
+    ]
+)
+class TestGradSimpleChainerX(GradTestBase, unittest.TestCase):
+
+    x_names = 'x',
+    y_names = 'y',
+
+    def forward(self):
+        self.y = self.x * self.x
+
+    def expected_grad(self):
+        grad = 2 * self.x * self.gy
+        return [grad]
+
+    def expected_double_grad(self):
+        ggrad = 2 * self.gy
+        return [ggrad]
+
+
 @testing.parameterize(*testing.product({
     'extend_graph_x': [False, True],
     'extend_graph_y': [False, True],
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
@@ -635,6 +635,12 @@ def setUp(self):
         self.optimizer = chainer.optimizers.SGD(lr)
 
     def test_update(self, backend_config):
+        if backend_config.xp is chainerx:
+            # ChainerX performs the loss scaling on its own backward
+            # method, the optimizer should not divide back the parameters
+            # This test is not actually creating a ChainerX
+            # computation graph so no actual loss scale is being done
+            self.optimizer.lr = 1.0
         target = self.target
         optimizer = self.optimizer
         target.to_device(backend_config.device)
diff --git a/tests/chainer_tests/test_variable.py b/tests/chainer_tests/test_variable.py
--- a/tests/chainer_tests/test_variable.py
+++ b/tests/chainer_tests/test_variable.py
@@ -2824,13 +2824,14 @@ def setUp(self):
         self.x = np.random.uniform(-1, 1, self.in_shape).astype(self.dtype)
         self.y = np.random.uniform(-1, 1, self.in_shape).astype(self.dtype)
 
-    def check_loss_scale(self, x_data, y_data):
+    def check_loss_scale(self, xp, x_data, y_data):
         x = chainer.Variable(x_data)
         y = chainer.Variable(y_data)
         z = x * y
         loss = F.sum(z)
         loss.backward(loss_scale=self.loss_scale)
-        if self.loss_scale is not None:
+        # ChainerX scales back gradients on the backward method
+        if xp is not chainerx and self.loss_scale is not None:
             x.grad /= self.loss_scale
             y.grad /= self.loss_scale
         rtol, atol = 1e-4, 1e-5
@@ -2840,11 +2841,24 @@ def check_loss_scale(self, x_data, y_data):
         testing.assert_allclose(y.data, x.grad, rtol=rtol, atol=atol)
 
     def test_loss_scale_cpu(self):
-        self.check_loss_scale(self.x, self.y)
+        self.check_loss_scale(np, self.x, self.y)
 
     @attr.gpu
     def test_loss_scale_gpu(self):
-        self.check_loss_scale(cuda.to_gpu(self.x), cuda.to_gpu(self.y))
+        self.check_loss_scale(cuda, cuda.to_gpu(self.x), cuda.to_gpu(self.y))
+
+    @attr.chainerx
+    def test_loss_scale_chainerx_cpu(self):
+        x = chainerx.array(self.x, device='native:0')
+        y = chainerx.array(self.y, device='native:0')
+        self.check_loss_scale(chainerx, x, y)
+
+    @attr.gpu
+    @attr.chainerx
+    def test_loss_scale_chainerx_gpu(self):
+        x = chainerx.array(self.x, device='cuda:0')
+        y = chainerx.array(self.y, device='cuda:0')
+        self.check_loss_scale(chainerx, x, y)
 
 
 @testing.parameterize(*testing.product({