Repository: chainer/chainer

Issue:
Support ChainerX in `optimizer_hooks.GradientClipping`
Part of #7634


Discussion:

Resolving patch:
diff --git a/chainer/optimizer_hooks/gradient_clipping.py b/chainer/optimizer_hooks/gradient_clipping.py
--- a/chainer/optimizer_hooks/gradient_clipping.py
+++ b/chainer/optimizer_hooks/gradient_clipping.py
@@ -5,22 +5,52 @@
 
 import chainer
 from chainer import backend
-from chainer import cuda
-
-
-def _sum_sqnorm(arr):
-    sq_sum = collections.defaultdict(float)
-    for x in arr:
-        with cuda.get_device_from_array(x) as dev:
-            x = x.ravel()
-            s = x.dot(x)
-            sq_sum[int(dev)] += s
-    # If only a single device is used, aggregate square norms on it.
-    if len(sq_sum) == 1:
-        with cuda.get_device_from_array(arr[0]):
-            return sum(six.itervalues(sq_sum))
+
+
+def _sum_sqnorm_grads(params):
+    # Calculates sum of squares of gradients.
+
+    # Returns a tuple of the sum and the device of the sum.
+    # The device will be `None` in multi-device case.
+
+    # If the inputs are on a single device, the sum is returned
+    # as an ndarray on the device, so that no synchronization is taken place.
+    # If there are multiple devices, accumulation is done on each device
+    # first, and the total sum is returned as a python float.
+
+    # TODO(niboshi): Support and test len(params) == 0
+
+    params_grouped = collections.defaultdict(list)
+    devices_map = {}
+
+    # Group params by devices.
+    for param in params:
+        device = param.device
+        params_grouped[device.name].append(param)
+        devices_map[device.name] = device
+
+    # Calculates partial sums for each device.
+    sq_sums = []
+    for device_name, paramlist in six.iteritems(params_grouped):
+        device = devices_map[device_name]
+        with chainer.using_device(device):
+            dots = []
+            for param in paramlist:
+                g = param.grad
+                g = g.ravel()
+                dots.append(g.dot(g))
+            sq_sums.append(sum(dots))
+
+    # Return the total sum.
+    if len(sq_sums) == 1:
+        # single device
+        sqnorm = sq_sums[0]
+        ret_device = params[0].device
     else:
-        return sum([float(i) for i in six.itervalues(sq_sum)])
+        # multi-device
+        sqnorm = sum([float(s) for s in sq_sums])
+        ret_device = None
+    return sqnorm, ret_device
 
 
 class GradientClipping(object):
@@ -52,19 +82,25 @@ def __init__(self, threshold):
         self.threshold = threshold
 
     def __call__(self, opt):
-        sqnorm = _sum_sqnorm([p.grad for p in opt.target.params(False)])
-        device = backend.get_device_from_array(sqnorm)
+        sqnorm, device = _sum_sqnorm_grads(list(opt.target.params(False)))
+        if device is None:
+            # Assign a dummy device for using_device.
+            device = backend.CpuDevice()
+
         with chainer.using_device(device):
             norm = device.xp.sqrt(sqnorm)
+            # TODO(niboshi): Could be inf if norm == 0
             rate = self.threshold / norm
-            # When no clipping is needed, skip the clipping on CPU and
-            # multiply 1.0 on the device otherwise.
+            # In NumPy backend, `rate` is already available on CPU and thus
+            # can be compared against 1 without extra overhead.
+            # Otherwise `clip` is used to avoid synchronization.
             if device.xp is numpy:
                 if rate >= 1:
                     return
             else:
                 rate = rate.clip(None, 1)
+
         for param in opt.target.params(False):
             grad = param.grad
-            with cuda.get_device_from_array(grad):
+            with chainer.using_device(param.device):
                 grad *= rate


Test patch:
diff --git a/tests/chainer_tests/optimizer_hooks_tests/test_gradient_clipping.py b/tests/chainer_tests/optimizer_hooks_tests/test_gradient_clipping.py
--- a/tests/chainer_tests/optimizer_hooks_tests/test_gradient_clipping.py
+++ b/tests/chainer_tests/optimizer_hooks_tests/test_gradient_clipping.py
@@ -1,110 +1,92 @@
+import math
 import unittest
 
 import numpy as np
 
 import chainer
-from chainer import backend
-from chainer import cuda
-import chainer.initializers as I
 from chainer import optimizer_hooks
-from chainer.optimizer_hooks import gradient_clipping
 from chainer import optimizers
 from chainer import testing
-from chainer.testing import attr
+
+
+_backend_params = [
+    # NumPy
+    {},
+    {'use_ideep': 'always'},
+    # CuPy
+    {'use_cuda': True, 'cuda_device': 0},
+    {'use_cuda': True, 'cuda_device': 1},
+    # ChainerX
+    {'use_chainerx': True, 'chainerx_device': 'native:0'},
+    {'use_chainerx': True, 'chainerx_device': 'cuda:0'},
+    {'use_chainerx': True, 'chainerx_device': 'cuda:1'},
+]
 
 
 class SimpleLink(chainer.Link):
 
-    def __init__(self, w, g):
+    def __init__(self, params):
         super(SimpleLink, self).__init__()
         with self.init_scope():
-            self.param = chainer.Parameter(I.Zero(), w.shape)
-        self.param.data = w
-        self.param.grad = g
-
-
-class TestOptimizerUtility(unittest.TestCase):
-
-    def setUp(self):
-        self.x = np.linspace(-1.0, 1.5, num=6).astype(np.float32).reshape(2, 3)
-        self.a = np.array(2.0)
-
-    def test_sqnorm_cpu(self):
-        # \Sum_{n=0}^{5} (-1.0+0.5n)**2 = 4.75
-        self.assertAlmostEqual(gradient_clipping._sum_sqnorm([self.x]), 4.75)
-
-    def test_sqnorm_scalar_cpu(self):
-        self.assertAlmostEqual(gradient_clipping._sum_sqnorm([self.a]), 4)
-
-    @attr.gpu
-    def test_sqnorm_gpu(self):
-        x = cuda.to_gpu(self.x)
-        self.assertAlmostEqual(gradient_clipping._sum_sqnorm([x]), 4.75)
-
-    @attr.gpu
-    def test_sqnorm_scalar_gpu(self):
-        a = cuda.to_gpu(self.a)
-        self.assertAlmostEqual(gradient_clipping._sum_sqnorm([a]), 4)
-
-    @attr.gpu
-    def test_sqnorm_array(self):
-        x = cuda.to_gpu(self.x)
-        a = cuda.to_gpu(self.a)
-        self.assertAlmostEqual(gradient_clipping._sum_sqnorm(
-            [self.x, self.a, x, a]), 8.75 * 2)
-
-    @attr.multi_gpu(2)
-    def test_sqnorm_array_multi_gpu(self):
-        x0 = cuda.to_gpu(self.x, device=0)
-        x1 = cuda.to_gpu(self.x, device=1)
-        a0 = cuda.to_gpu(self.a, device=0)
-        a1 = cuda.to_gpu(self.a, device=1)
-        self.assertAlmostEqual(gradient_clipping._sum_sqnorm(
-            [self.x, self.a, x0, a0, x1, a1]), 8.75 * 3)
+            for i, p in enumerate(params):
+                setattr(self, 'p{}'.format(i), p)
 
 
+@testing.backend.inject_backend_tests(None, _backend_params)
+@testing.backend.inject_backend_tests(None, _backend_params)
+@testing.backend.inject_backend_tests(None, _backend_params)
 class TestGradientClipping(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_clipping(self, multiplier):
-        w = self.target.param.data
-        g = self.target.param.grad
-        xp = backend.get_array_module(w)
-
-        norm = xp.sqrt(gradient_clipping._sum_sqnorm(g))
-        threshold = norm * multiplier
-        if multiplier < 1:
-            expect = w - g * multiplier
-        else:
-            expect = w - g
+        num_params = 3
+        arrs = [
+            np.random.uniform(-3, 3, (2, 3)).astype(np.float32)
+            for _ in range(num_params)]
+        grads = [
+            np.random.uniform(-3, 3, (2, 3)).astype(np.float32)
+            for _ in range(num_params)]
+        params = []
+        for arr, grad in zip(arrs, grads):
+            param = chainer.Parameter(arr)
+            param.grad = grad
+            params.append(param)
+
+        self.target = SimpleLink(params)
+        self.norm = math.sqrt(sum([g.ravel().dot(g.ravel()) for g in grads]))
+
+    def check_clipping(self, backend_configs, rate):
+        target = self.target
+        norm = self.norm
+        assert len(backend_configs) == len(list(target.params()))
+        devices = [bc.device for bc in backend_configs]
+
+        threshold = norm * rate
+
+        expects = []
+        for param, device in zip(target.params(), devices):
+            expects.append(param.array - param.grad * min(1, rate))
+            param.to_device(device)
 
         opt = optimizers.SGD(lr=1)
-        opt.setup(self.target)
+        opt.setup(target)
         opt.add_hook(
             optimizer_hooks.GradientClipping(threshold))
         opt.update()
 
-        testing.assert_allclose(expect, w)
-
-    def test_clipping_cpu(self):
-        self.check_clipping(0.5)
-
-    @attr.gpu
-    def test_clipping_gpu(self):
-        self.target.to_gpu()
-        self.check_clipping(0.5)
+        for expect, param in zip(expects, target.params()):
+            testing.assert_allclose(expect, param.array)
 
-    def test_clipping_2_cpu(self):
-        self.check_clipping(2.0)
+    def test_clipping(
+            self, backend_config0, backend_config1, backend_config2):
+        self.check_clipping(
+            [backend_config0, backend_config1, backend_config2],
+            0.5)
 
-    @attr.gpu
-    def test_clipping_2_gpu(self):
-        self.target.to_gpu()
-        self.check_clipping(2.0)
+    def test_clipping_2(
+            self, backend_config0, backend_config1, backend_config2):
+        self.check_clipping(
+            [backend_config0, backend_config1, backend_config2],
+            2.0)
 
 
 testing.run_module(__name__, __file__)