Repository: chainer/chainer

Issue:
Adam optimizer's default eps is too small for FP16
Rel #6168.

Currently, Adam optimizer's defaut eps is `1e-8`, which is too small for a FP16 parameter to update and causes a `ValueError` without an explicitly supplied proper eps.

Some options are:
- Make `Adam` take `dtype` parameter to initialize eps properly. If the `dtype` is not supplied, the default dtype is used.
- Lazily initialize the eps on updating for the dtype of the parameter to update
- (Explicitly pass proper eps to `Adam` as it is)



Discussion:

Resolving patch:
diff --git a/chainer/optimizers/adam.py b/chainer/optimizers/adam.py
--- a/chainer/optimizers/adam.py
+++ b/chainer/optimizers/adam.py
@@ -48,6 +48,26 @@ def _learning_rate(hp, t):
     return hp.alpha * math.sqrt(fix2) / fix1
 
 
+def _get_intermediate_dtype(dtype):
+    # Returns the dtype for intermediate calculation.
+    # For float16 input, float32 is used.
+    # Otherwise the same dtype as the parameter is used.
+    if dtype == numpy.float16:
+        return numpy.float32
+    return dtype
+
+
+def _inplace_axpby(x, a, b, y):
+    # in-place axpby: x = a * x + b * y
+    if isinstance(x, intel64.mdarray):
+        x.inplace_axpby(a, b, y)
+    else:
+        if a == 1:
+            x += b * y
+        else:
+            x[:] = a * x + b * y
+
+
 class AdamRule(optimizer.UpdateRule):
 
     """Update rule of Adam optimization algorithm.
@@ -83,6 +103,12 @@ class AdamRule(optimizer.UpdateRule):
     _kernel = None
     _amsgrad_kernel = None
 
+    # Only used in `update_core_gpu`.
+    # A dummy ndarray to help ElementwiseKernel deduce generic type T as
+    # `dtype`.
+    # It cannot be deduced only by scalar arguments.
+    _dummy = None
+
     def __init__(self, parent_hyperparam=None,
                  alpha=None, beta1=None, beta2=None, eps=None,
                  eta=None, weight_decay_rate=None, amsgrad=None):
@@ -118,60 +144,69 @@ def init_state(self, param):
             self.state['v'] = intel64.ideep.array(
                 self.state['v'], itype=intel64.ideep.wgt_array)
 
+    def _check_eps(self, interm_dtype):
+        # Checks that the eps does not underflow.
+        hp = self.hyperparam
+        eps = interm_dtype(hp.eps)
+        if hp.eps != 0 and eps == 0:
+            raise ValueError(
+                'eps of Adam optimizer is too small for {} ({})'.format(
+                    interm_dtype.name, hp.eps))
+        # Note that the converted `eps` (numpy scalar) is discarded here and
+        # the original `hp.eps` is used in calculation, because Python
+        # scalars are faster in cupy elementwise kernels.
+
     def update_core_cpu(self, param):
         grad = param.grad
         if grad is None:
             return
         hp = self.hyperparam
-        eps = grad.dtype.type(hp.eps)
-        if hp.eps != 0 and eps == 0:
-            raise ValueError(
-                'eps of Adam optimizer is too small for {} ({})'.format(
-                    grad.dtype.name, hp.eps))
+        dtype = _get_intermediate_dtype(param.dtype.type)
+        self._check_eps(dtype)
+        grad = grad.astype(dtype, copy=False)
+
         m, v = self.state['m'], self.state['v']
-        if (isinstance(m, intel64.mdarray)
-                and isinstance(v, intel64.mdarray)):
-            m.inplace_axpby(1.0, 1.0 - hp.beta1, grad - m)
-            v.inplace_axpby(1.0, 1.0 - hp.beta2, grad*grad - v)
-            if hp.amsgrad:
-                vhat = self.state['vhat']
-                numpy.maximum(vhat, v, out=vhat)
-            else:
-                vhat = v
-            param.data.inplace_axpby(
-                1.0 - hp.weight_decay_rate, -hp.eta,
-                self.alpha_t * m / (numpy.sqrt(vhat) + hp.eps))
+
+        # m += (1 - beta1) * (grad - m)
+        _inplace_axpby(m, 1.0, 1.0 - hp.beta1, grad - m)
+        # v += (1 - beta2) * (grad * grad - v)
+        _inplace_axpby(v, 1.0, 1.0 - hp.beta2, grad*grad - v)
+
+        if hp.amsgrad:
+            vhat = self.state['vhat']
+            numpy.maximum(vhat, v, out=vhat)
         else:
-            m += (1 - hp.beta1) * (grad - m)
-            v += (1 - hp.beta2) * (grad * grad - v)
-            if hp.amsgrad:
-                vhat = self.state['vhat']
-                numpy.maximum(vhat, v, out=vhat)
-            else:
-                vhat = v
-            param.data -= hp.eta * (
-                self.alpha_t * m / (numpy.sqrt(vhat) + hp.eps) +
-                hp.weight_decay_rate * param.data)
+            vhat = v
+        vhat = vhat.astype(dtype, copy=False)
+
+        # param -=
+        #  eta * (alpha_t * m / (sqrt(vhat) + eps) - weight_decay_rate * param)
+        _inplace_axpby(
+            param.data,
+            1.0 - hp.weight_decay_rate,
+            -hp.eta,
+            self.alpha_t * m / (numpy.sqrt(vhat) + hp.eps))
 
     def update_core_gpu(self, param):
         grad = param.grad
         if grad is None:
             return
-
         hp = self.hyperparam
-        eps = grad.dtype.type(hp.eps)
-        if hp.eps != 0 and eps == 0:
-            raise ValueError(
-                'eps of Adam optimizer is too small for {} ({})'.format(
-                    grad.dtype.name, hp.eps))
+        dtype = _get_intermediate_dtype(param.dtype.type)
+        self._check_eps(dtype)
+
+        if self._dummy is None:
+            self._dummy = cuda.cupy.empty((0,), dtype=dtype)
+
         if hp.amsgrad:
             if AdamRule._amsgrad_kernel is None:
                 AdamRule._amsgrad_kernel = cuda.elementwise(
-                    'T grad, T alpha_t, T one_minus_beta1, T one_minus_beta2, '
-                    'T eps, T eta, T weight_decay_rate',
-                    'T param, T m, T v, T vhat',
-                    '''m += one_minus_beta1 * (grad - m);
-                       v += one_minus_beta2 * (grad * grad - v);
+                    'P grad, T alpha_t, T one_minus_beta1, T one_minus_beta2, '
+                    'T eps, T eta, T weight_decay_rate, raw T dummy',
+                    'P param, P m, P v, P vhat',
+                    '''T grad_ = static_cast<T>(grad);
+                       m += one_minus_beta1 * (grad_ - m);
+                       v += one_minus_beta2 * (grad_ * grad_ - v);
                        vhat = max(vhat, v);
                        param -= eta * (alpha_t * m / (sqrt(vhat) + eps) +
                                        weight_decay_rate * param);''',
@@ -179,24 +214,26 @@ def update_core_gpu(self, param):
             AdamRule._amsgrad_kernel(
                 grad, self.alpha_t, 1 - hp.beta1,
                 1 - hp.beta2, hp.eps,
-                hp.eta, hp.weight_decay_rate,
+                hp.eta, hp.weight_decay_rate, self._dummy,
                 param.data, self.state['m'], self.state['v'],
                 self.state['vhat'])
         else:
             if AdamRule._kernel is None:
                 AdamRule._kernel = cuda.elementwise(
-                    'T grad, T alpha_t, T one_minus_beta1, T one_minus_beta2, '
-                    'T eps, T eta, T weight_decay_rate',
-                    'T param, T m, T v',
-                    '''m += one_minus_beta1 * (grad - m);
-                       v += one_minus_beta2 * (grad * grad - v);
+                    'P grad, T alpha_t, T one_minus_beta1, T one_minus_beta2, '
+                    'T eps, T eta, T weight_decay_rate, raw T dummy',
+                    'P param, P m, P v',
+                    '''T grad_ = static_cast<T>(grad);
+                       m += one_minus_beta1 * (grad_ - m);
+                       v += one_minus_beta2 * (grad_ * grad_ - v);
                        param -= eta * (alpha_t * m / (sqrt(v) + eps) +
                                        weight_decay_rate * param);''',
                     'adam')
-            AdamRule._kernel(grad, self.alpha_t, 1 - hp.beta1,
-                             1 - hp.beta2, hp.eps,
-                             hp.eta, hp.weight_decay_rate,
-                             param.data, self.state['m'], self.state['v'])
+            AdamRule._kernel(
+                grad, self.alpha_t, 1 - hp.beta1,
+                1 - hp.beta2, hp.eps,
+                hp.eta, hp.weight_decay_rate, self._dummy,
+                param.data, self.state['m'], self.state['v'])
 
     @property
     def alpha_t(self):


Test patch:
diff --git a/tests/chainer_tests/optimizers_tests/test_optimizers_by_linear_model.py b/tests/chainer_tests/optimizers_tests/test_optimizers_by_linear_model.py
--- a/tests/chainer_tests/optimizers_tests/test_optimizers_by_linear_model.py
+++ b/tests/chainer_tests/optimizers_tests/test_optimizers_by_linear_model.py
@@ -192,11 +192,7 @@ def create(self):
 class TestAdam(OptimizerTestBase, unittest.TestCase):
 
     def create(self):
-        if self.dtype == numpy.float16:
-            kwargs = {'eps': 1e-6}
-        else:
-            kwargs = {}
-        return optimizers.Adam(0.05, **kwargs)
+        return optimizers.Adam(0.05)
 
 
 @testing.parameterize(*testing.product({