Repository: chainer/chainer

Issue:
GroupNormalization fails with CHAINER_DTYPE=float16 or CHAINER_DTYPE=mixed16
```
% cat main.py                         
import numpy

import chainer
import chainer.links
import cupy


gn = chainer.links.GroupNormalization(1)
gn.to_gpu()

dtype = chainer.get_dtype()
x = chainer.Variable(cupy.ones((1, 2, 3, 4), dtype=dtype))
y = gn(x)

y.grad = cupy.ones_like(y.array)
y.backward()
```
```
% CHAINER_DTYPE=float32 python main.py
```
```
% CHAINER_DTYPE=float16 python main.py
Traceback (most recent call last):
  File "main.py", line 16, in <module>
    y.backward()
  File "/path/to/chainer/chainer/variable.py", line 1470, in backward
    [(node, grad_var)], retain_grad, loss_scale)
  File "/path/to/chainer/chainer/_backprop.py", line 226, in _backprop_to_all
    func, target_input_indexes, out_grad, in_grad, is_debug)
  File "/path/to/chainer/chainer/_backprop_utils.py", line 140, in backprop_step
    _reraise_with_stack(func, e)
  File "/path/to/chainer/chainer/_backprop_utils.py", line 138, in backprop_step
    target_input_indexes, grad_outputs)
  File "/path/to/chainer/chainer/functions/normalization/group_normalization.py", line 122, in backward
    self.dummy_gamma).apply((x,))
  File "/path/to/chainer/chainer/function_node.py", line 325, in apply
    outputs = self.forward(in_data)
  File "/path/to/chainer/chainer/function_node.py", line 518, in forward
    return self.forward_gpu(inputs)
  File "/path/to/chainer/chainer/functions/normalization/group_normalization.py", line 201, in forward_gpu
    'groupnorm_x_hat')(x, self.mean[:, None], self.inv_std[:, None])
  File "cupy/core/_kernel.pyx", line 539, in cupy.core._kernel.ElementwiseKernel.__call__
  File "cupy/core/_kernel.pyx", line 582, in cupy.core._kernel.ElementwiseKernel._decide_params_type
  File "cupy/core/_kernel.pyx", line 292, in cupy.core._kernel._decide_params_type_core
TypeError: Type is mismatched. mean <class 'numpy.float32'> <class 'numpy.float16'> T
```
```
% CHAINER_DTYPE=mixed16 python main.py
Traceback (most recent call last):
  File "main.py", line 16, in <module>
    y.backward()
  File "/path/to/chainer/chainer/variable.py", line 1470, in backward
    [(node, grad_var)], retain_grad, loss_scale)
  File "/path/to/chainer/chainer/_backprop.py", line 226, in _backprop_to_all
    func, target_input_indexes, out_grad, in_grad, is_debug)
  File "/path/to/chainer/chainer/_backprop_utils.py", line 140, in backprop_step
    _reraise_with_stack(func, e)
  File "/path/to/chainer/chainer/_backprop_utils.py", line 138, in backprop_step
    target_input_indexes, grad_outputs)
  File "/path/to/chainer/chainer/functions/normalization/group_normalization.py", line 122, in backward
    self.dummy_gamma).apply((x,))
  File "/path/to/chainer/chainer/function_node.py", line 325, in apply
    outputs = self.forward(in_data)
  File "/path/to/chainer/chainer/function_node.py", line 518, in forward
    return self.forward_gpu(inputs)
  File "/path/to/chainer/chainer/functions/normalization/group_normalization.py", line 201, in forward_gpu
    'groupnorm_x_hat')(x, self.mean[:, None], self.inv_std[:, None])
  File "cupy/core/_kernel.pyx", line 539, in cupy.core._kernel.ElementwiseKernel.__call__
  File "cupy/core/_kernel.pyx", line 582, in cupy.core._kernel.ElementwiseKernel._decide_params_type
  File "cupy/core/_kernel.pyx", line 292, in cupy.core._kernel._decide_params_type_core
TypeError: Type is mismatched. mean <class 'numpy.float32'> <class 'numpy.float16'> T
```


Discussion:

Resolving patch:
diff --git a/chainer/functions/normalization/group_normalization.py b/chainer/functions/normalization/group_normalization.py
--- a/chainer/functions/normalization/group_normalization.py
+++ b/chainer/functions/normalization/group_normalization.py
@@ -1,3 +1,5 @@
+import numpy
+
 import chainer
 from chainer import backend
 from chainer.backends import cuda
@@ -31,8 +33,8 @@ def check_type_forward(self, in_types):
             x_type.ndim >= 2,
             gamma_type.ndim == 1,
             beta_type.ndim == 1,
-            gamma_type.dtype == x_type.dtype,
-            beta_type.dtype == x_type.dtype,
+            gamma_type.dtype.kind == 'f',
+            gamma_type.dtype == beta_type.dtype,
             x_type.shape[1] == gamma_type.shape[0],
             gamma_type.shape == beta_type.shape,
         )
@@ -50,19 +52,24 @@ def forward(self, inputs):
         self.retain_inputs((0, 1))
         x, gamma, beta = inputs
 
+        interm_dtype = numpy.promote_types(x.dtype, gamma.dtype)
+
+        gamma = gamma.astype(interm_dtype, copy=False)
+        beta = beta.astype(interm_dtype, copy=False)
+
         orig_shape = x.shape
         batch_size, channels = orig_shape[:2]
         groups = self.groups
         reduced_shape = (batch_size * groups, -1)
         x = x.reshape(reduced_shape)
 
-        self.mean = x.mean(axis=1)
+        self.mean = x.mean(axis=1, dtype=interm_dtype)
         x_hat = x - self.mean[:, None]
         var = (x_hat * x_hat).mean(axis=1)
         var += self.eps
         self.inv_std = var
         del var
-        xp.sqrt(self.inv_std, out=self.inv_std, dtype=x.dtype)
+        xp.sqrt(self.inv_std, out=self.inv_std)
         xp.reciprocal(self.inv_std, out=self.inv_std)
         x_hat *= self.inv_std[:, None]
 
@@ -71,7 +78,7 @@ def forward(self, inputs):
         y += beta[:, None]
 
         y = y.reshape(orig_shape)
-        return y,
+        return y.astype(x.dtype, copy=False),
 
     def forward_cudnn(self, inputs):
         if self.eps < libcudnn.CUDNN_BN_MIN_EPSILON:
@@ -83,6 +90,11 @@ def forward_cudnn(self, inputs):
         x, gamma, beta = inputs
         xp = cuda.cupy
 
+        interm_dtype = numpy.promote_types(x.dtype, gamma.dtype)
+
+        gamma = gamma.astype(interm_dtype, copy=False)
+        beta = beta.astype(interm_dtype, copy=False)
+
         orig_shape = x.shape
         batch_size, channels = orig_shape[:2]
         groups = self.groups
@@ -90,7 +102,7 @@ def forward_cudnn(self, inputs):
         x = x.reshape(cudnn_shape)
 
         with x.device:
-            dummy_beta = xp.zeros(batch_size * groups, dtype=x.dtype)
+            dummy_beta = xp.zeros(batch_size * groups, dtype=beta.dtype)
             self.dummy_gamma = xp.ones_like(dummy_beta)
         x_hat, self.mean, self.inv_std = \
             cudnn.batch_normalization_forward_training(
@@ -100,7 +112,7 @@ def forward_cudnn(self, inputs):
 
         y = x_hat.reshape((batch_size, channels, -1))
         cuda.elementwise(
-            'T gamma, T beta', 'T y',
+            'T gamma, T beta', 'U y',
             'y = y * gamma + beta',
             'groupnorm_y')(gamma[:, None], beta[:, None], y)
 
@@ -111,22 +123,29 @@ def backward(self, indexes, grad_outputs):
         x, gamma = self.get_retained_inputs()
         gy, = grad_outputs
 
+        interm_dtype = numpy.promote_types(x.dtype, gamma.dtype)
+
+        gamma = chainer.functions.cast(gamma, interm_dtype)
+
         orig_shape = x.shape
         batch_size = orig_shape[0]
         groups = self.groups
         reduced_shape = (batch_size * groups, -1)
         x = x.reshape(reduced_shape)
+        x_ = chainer.functions.cast(x, interm_dtype)
 
         x_hat, = _XHat(
             self.eps, self.mean, self.inv_std,
-            self.dummy_gamma).apply((x,))
-        gx_hat, ggamma, gbeta = _ScaleShiftGrad().apply((x_hat, gamma, gy))
+            self.dummy_gamma).apply((x_,))
+        gx_hat, ggamma, gbeta = _ScaleShiftGrad().apply((
+            x_hat, gamma, chainer.functions.cast(gy, interm_dtype)))
         gx, = _XHatGrad(
             self.eps, self.mean, self.inv_std,
-            self.dummy_gamma, x_hat.array).apply((x, gx_hat))
+            self.dummy_gamma, x_hat.array).apply(
+                (x_, gx_hat))
 
         gx = gx.reshape(orig_shape)
-        return gx, ggamma, gbeta
+        return chainer.functions.cast(gx, x.dtype), ggamma, gbeta
 
 
 class _ScaleShiftGrad(function_node.FunctionNode):
@@ -196,7 +215,7 @@ def forward_gpu(self, inputs):
         self.retain_inputs((0,))
         x, = inputs
         x_hat = cuda.elementwise(
-            'T x, T mean, T inv_std', 'T x_hat',
+            'T x, U mean, U inv_std', 'T x_hat',
             'x_hat = (x - mean) * inv_std',
             'groupnorm_x_hat')(x, self.mean[:, None], self.inv_std[:, None])
         self.retain_outputs((0,))
diff --git a/chainer/links/normalization/group_normalization.py b/chainer/links/normalization/group_normalization.py
--- a/chainer/links/normalization/group_normalization.py
+++ b/chainer/links/normalization/group_normalization.py
@@ -1,4 +1,8 @@
+import numpy
+
+import chainer
 from chainer.functions.normalization import group_normalization
+from chainer import initializers
 from chainer import link
 from chainer import variable
 
@@ -49,10 +53,19 @@ def __init__(self, groups, size=None, eps=1e-5, initial_gamma=None,
         if initial_beta is None:
             initial_beta = 0
 
+        highprec_dtype = chainer.get_dtype(
+            None, map_mixed16=numpy.float32)
+
         with self.init_scope():
             self.groups = groups
-            self.gamma = variable.Parameter(initial_gamma)
-            self.beta = variable.Parameter(initial_beta)
+            gamma_initializer = \
+                initializers._get_initializer(initial_gamma)
+            gamma_initializer.dtype = highprec_dtype
+            beta_initializer = \
+                initializers._get_initializer(initial_beta)
+            beta_initializer.dtype = highprec_dtype
+            self.gamma = variable.Parameter(gamma_initializer)
+            self.beta = variable.Parameter(beta_initializer)
             self.eps = eps
 
         if size is not None:


Test patch:
diff --git a/tests/chainer_tests/links_tests/normalization_tests/test_group_normalization.py b/tests/chainer_tests/links_tests/normalization_tests/test_group_normalization.py
--- a/tests/chainer_tests/links_tests/normalization_tests/test_group_normalization.py
+++ b/tests/chainer_tests/links_tests/normalization_tests/test_group_normalization.py
@@ -14,19 +14,25 @@
 @testing.parameterize(*(testing.product({
     'shape': [(1, 4, 5, 3), (5, 4, 7), (3, 20)],
     'groups': [1, 2, 4],
-    'dtype': [numpy.float32],
+    'dtype': [numpy.float16, numpy.float32, numpy.float64,
+              chainer.mixed16],
 })))
 class GroupNormalizationTest(unittest.TestCase):
 
     def setUp(self):
-        self.link = links.GroupNormalization(self.groups)
+        with chainer.using_config('dtype', self.dtype):
+            self.link = links.GroupNormalization(self.groups)
         self.link.cleargrads()
 
         self.x, = self.generate_inputs()
         self.gy = numpy.random.uniform(-1, 1, self.shape).astype(self.dtype)
 
-        self.check_forward_options = {'atol': 1e-4, 'rtol': 1e-3}
-        self.check_backward_options = {'atol': 1e-3, 'rtol': 1e-2}
+        if self.dtype in (numpy.float16, chainer.mixed16):
+            self.check_forward_options = {'atol': 1e-2, 'rtol': 1e-1}
+            self.check_backward_options = {'atol': 5e-1, 'rtol': 1e-1}
+        else:
+            self.check_forward_options = {'atol': 1e-4, 'rtol': 1e-3}
+            self.check_backward_options = {'atol': 1e-3, 'rtol': 1e-2}
 
     def generate_inputs(self):
         shape = self.shape
@@ -82,82 +88,106 @@ def check_backward(self, x_data, y_grad):
         gradient_check.check_backward(
             self.link, x_data, y_grad,
             (self.link.gamma, self.link.beta),
-            eps=1e-2, **self.check_backward_options)
+            eps=2e-2, **self.check_backward_options)
 
     def test_backward_cpu(self):
-        self.link(numpy.zeros(self.shape, dtype='f'))
+        self.link(numpy.zeros(self.shape, dtype=self.dtype))
         self.check_backward(self.x, self.gy)
 
     @attr.gpu
     def test_backward_gpu(self):
         self.link.to_gpu()
-        self.link(cuda.cupy.zeros(self.shape, dtype='f'))
+        self.link(cuda.cupy.zeros(self.shape, dtype=self.dtype))
         self.check_backward(cuda.to_gpu(self.x), cuda.to_gpu(self.gy))
 
     @attr.cudnn
     def test_backward_gpu_without_cudnn(self):
         self.link.use_cudnn = False
-        self.link(numpy.zeros(self.shape, dtype='f'))
+        self.link(numpy.zeros(self.shape, dtype=self.dtype))
         self.test_backward_gpu()
 
 
 @testing.parameterize(*testing.product({
     'size': [3, 30],
-    'groups': [1, 3]
+    'groups': [1, 3],
+    'dtype': [numpy.float16, numpy.float32, numpy.float64,
+              chainer.mixed16],
 }))
 class TestInitialize(unittest.TestCase):
 
     def setUp(self):
+        self.lowprec_dtype = chainer.get_dtype(self.dtype)
+        highprec_dtype = chainer.get_dtype(
+            self.dtype, map_mixed16=numpy.float32)
+
         self.initial_gamma = numpy.random.uniform(-1, 1, self.size)
-        self.initial_gamma = self.initial_gamma.astype(numpy.float32)
+        self.initial_gamma = self.initial_gamma.astype(highprec_dtype)
         self.initial_beta = numpy.random.uniform(-1, 1, self.size)
-        self.initial_beta = self.initial_beta.astype(numpy.float32)
+        self.initial_beta = self.initial_beta.astype(highprec_dtype)
         self.link = links.GroupNormalization(self.groups,
                                              initial_gamma=self.initial_gamma,
                                              initial_beta=self.initial_beta)
         self.shape = (1, self.size, 1)
 
     def test_initialize_cpu(self):
-        self.link(numpy.zeros(self.shape, dtype='f'))
+        with chainer.using_config('dtype', self.dtype):
+            self.link(numpy.zeros(self.shape, dtype=self.lowprec_dtype))
         testing.assert_allclose(self.initial_gamma, self.link.gamma.data)
         testing.assert_allclose(self.initial_beta, self.link.beta.data)
 
     @attr.gpu
     def test_initialize_gpu(self):
         self.link.to_gpu()
-        self.link(cuda.cupy.zeros(self.shape, dtype='f'))
+        with chainer.using_config('dtype', self.dtype):
+            self.link(cuda.cupy.zeros(self.shape, dtype=self.lowprec_dtype))
         testing.assert_allclose(self.initial_gamma, self.link.gamma.data)
         testing.assert_allclose(self.initial_beta, self.link.beta.data)
 
 
 @testing.parameterize(*testing.product({
     'size': [3, 30],
-    'groups': [1, 3]
+    'groups': [1, 3],
+    'dtype': [numpy.float16, numpy.float32, numpy.float64,
+              chainer.mixed16],
 }))
 class TestDefaultInitializer(unittest.TestCase):
 
     def setUp(self):
+        self.lowprec_dtype = chainer.get_dtype(self.dtype)
+        self.highprec_dtype = chainer.get_dtype(
+            self.dtype, map_mixed16=numpy.float32)
+
         self.size = 3
-        self.link = links.GroupNormalization(self.groups)
+        with chainer.using_config('dtype', self.dtype):
+            self.link = links.GroupNormalization(self.groups)
         self.shape = (1, self.size, 1)
 
     def test_initialize_cpu(self):
-        self.link(numpy.zeros(self.shape, dtype='f'))
+        self.link(numpy.zeros(self.shape, dtype=self.lowprec_dtype))
+
         testing.assert_allclose(numpy.ones(self.size), self.link.gamma.data)
+        self.assertEqual(self.link.gamma.dtype, self.highprec_dtype)
+
         testing.assert_allclose(
             numpy.zeros(self.size), self.link.beta.data)
+        self.assertEqual(self.link.beta.dtype, self.highprec_dtype)
 
     @attr.gpu
     def test_initialize_gpu(self):
         self.link.to_gpu()
-        self.link(cuda.cupy.zeros(self.shape, dtype='f'))
+        self.link(cuda.cupy.zeros(self.shape, dtype=self.lowprec_dtype))
+
         testing.assert_allclose(numpy.ones(self.size), self.link.gamma.data)
+        self.assertEqual(self.link.gamma.dtype, self.highprec_dtype)
+
         testing.assert_allclose(
             numpy.zeros(self.size), self.link.beta.data)
+        self.assertEqual(self.link.beta.dtype, self.highprec_dtype)
 
 
 @testing.parameterize(*testing.product({
     'shape': [(2,), ()],
+    'dtype': [numpy.float16, numpy.float32, numpy.float64],
 }))
 class TestInvalidInput(unittest.TestCase):
 
@@ -166,20 +196,26 @@ def setUp(self):
 
     def test_invalid_shape_cpu(self):
         with self.assertRaises(ValueError):
-            self.link(chainer.Variable(numpy.zeros(self.shape, dtype='f')))
+            self.link(chainer.Variable(
+                numpy.zeros(self.shape, dtype=self.dtype)))
 
     @attr.gpu
     def test_invalid_shape_gpu(self):
         self.link.to_gpu()
         with self.assertRaises(ValueError):
-            self.link(chainer.Variable(cuda.cupy.zeros(self.shape, dtype='f')))
+            self.link(
+                chainer.Variable(
+                    cuda.cupy.zeros(self.shape, dtype=self.dtype)))
 
 
+@testing.parameterize(*testing.product({
+    'dtype': [numpy.float16, numpy.float32, numpy.float64],
+}))
 class TestInvalidInitialize(unittest.TestCase):
 
     def setUp(self):
         shape = (2, 5, 2)
-        self.x = chainer.Variable(numpy.zeros(shape, dtype='f'))
+        self.x = chainer.Variable(numpy.zeros(shape, dtype=self.dtype))
 
     def test_invalid_groups(self):
         self.link = links.GroupNormalization(groups=3)