You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@mxnet.apache.org by jx...@apache.org on 2017/10/08 05:22:50 UTC
[incubator-mxnet] branch master updated: Making mixed precision
work with all optimizers (#7654)
This is an automated email from the ASF dual-hosted git repository.
jxie pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-mxnet.git
The following commit(s) were added to refs/heads/master by this push:
new ee97715 Making mixed precision work with all optimizers (#7654)
ee97715 is described below
commit ee9771580765393b194172bab169dd67ac575209
Author: Przemyslaw Tredak <pt...@gmail.com>
AuthorDate: Sat Oct 7 22:22:44 2017 -0700
Making mixed precision work with all optimizers (#7654)
* Making mixed precision work with all optimizers
* Restart CI
* Restart CI
---
example/image-classification/common/fit.py | 7 +-
python/mxnet/optimizer.py | 103 ++++++++++++++++++----
tests/python/unittest/test_optimizer.py | 132 ++++++++++++++++-------------
3 files changed, 160 insertions(+), 82 deletions(-)
diff --git a/example/image-classification/common/fit.py b/example/image-classification/common/fit.py
index aeead0f..51a1abe 100755
--- a/example/image-classification/common/fit.py
+++ b/example/image-classification/common/fit.py
@@ -164,11 +164,8 @@ def fit(args, network, data_loader, **kwargs):
optimizer_params = {
'learning_rate': lr,
'wd' : args.wd,
- 'lr_scheduler': lr_scheduler}
-
- # Add 'multi_precision' parameter only for SGD optimizer
- if args.optimizer == 'sgd':
- optimizer_params['multi_precision'] = True
+ 'lr_scheduler': lr_scheduler,
+ 'multi_precision': True}
# Only a limited number of optimizers have 'momentum' property
has_momentum = {'sgd', 'dcasgd', 'nag'}
diff --git a/python/mxnet/optimizer.py b/python/mxnet/optimizer.py
index 9f89415..4aefb66 100644
--- a/python/mxnet/optimizer.py
+++ b/python/mxnet/optimizer.py
@@ -24,7 +24,7 @@ import logging
import warnings
import numpy
from .base import py_str
-from .ndarray import (NDArray, zeros, clip, sqrt, array, maximum, abs as NDabs)
+from .ndarray import (NDArray, zeros, clip, sqrt, cast, maximum, abs as NDabs)
from .ndarray import (sgd_update, sgd_mom_update, adam_update, rmsprop_update, rmspropalex_update,
mp_sgd_update, mp_sgd_mom_update, ftrl_update)
from .random import normal
@@ -61,6 +61,12 @@ class Optimizer(object):
begin_num_update : int, optional
The initial number of updates.
+ multi_precision : bool, optional
+ Flag to control the internal precision of the optimizer.
+ ``False`` results in using the same precision as the weights (default),
+ ``True`` makes internal 32-bit copy of the weights and applies gradients
+ in 32-bit precision even if actual weights used in the model have lower precision.
+ Turning this on can improve convergence and accuracy when training with float16.
Properties
----------
@@ -71,7 +77,7 @@ class Optimizer(object):
def __init__(self, rescale_grad=1., param_idx2name=None, wd=0.,
clip_gradient=None, learning_rate=0.01,
lr_scheduler=None, sym=None, begin_num_update=0,
- param_dict=None):
+ multi_precision=False, param_dict=None):
self.rescale_grad = rescale_grad
self.lr = learning_rate
self.lr_scheduler = lr_scheduler
@@ -85,6 +91,7 @@ class Optimizer(object):
self.num_update = begin_num_update
self._index_update_count = {}
self.clip_gradient = clip_gradient
+ self.multi_precision = multi_precision
if param_idx2name is None:
param_idx2name = {}
@@ -189,6 +196,36 @@ class Optimizer(object):
The state associated with the weight.
"""
+ def create_state_multi_precision(self, index, weight):
+ """Creates auxiliary state for a given weight, including FP32 high
+ precision copy if original weight is FP16.
+
+ This method is provided to perform automatic mixed precision training
+ for optimizers that do not support it themselves.
+
+ Parameters
+ ----------
+ index : int
+ An unique index to identify the weight.
+ weight : NDArray
+ The weight.
+
+ Returns
+ -------
+ state : any obj
+ The state associated with the weight.
+ """
+ weight_master_copy = None
+ if self.multi_precision and weight.dtype == numpy.float16:
+ weight_master_copy = weight.astype(numpy.float32)
+ return (weight_master_copy,) + (self.create_state(index, weight_master_copy),)
+ if weight.dtype == numpy.float16 and not self.multi_precision:
+ warnings.warn("Accumulating with float16 in optimizer can lead to "
+ "poor accuracy or slow convergence. "
+ "Consider using multi_precision=True option of the "
+ "optimizer")
+ return self.create_state(index, weight)
+
def update(self, index, weight, grad, state):
"""Updates the given parameter using the corresponding gradient and state.
@@ -207,6 +244,33 @@ class Optimizer(object):
"""
raise NotImplementedError()
+ def update_multi_precision(self, index, weight, grad, state):
+ """Updates the given parameter using the corresponding gradient and state.
+ Mixed precision version.
+
+ Parameters
+ ----------
+ index : int
+ The unique index of the parameter into the individual learning
+ rates and weight decays. Learning rates and weight decay
+ may be set via `set_lr_mult()` and `set_wd_mult()`, respectively.
+ weight : NDArray
+ The parameter to be updated.
+ grad : NDArray
+ The gradient of the objective with respect to this parameter.
+ state : any obj
+ The state returned by `create_state()`.
+ """
+ if self.multi_precision and weight.dtype == numpy.float16:
+ # Wrapper for mixed precision
+ weight_master_copy = state[0]
+ original_state = state[1]
+ grad32 = grad.astype(numpy.float32)
+ self.update(index, weight_master_copy, grad32, original_state)
+ cast(weight_master_copy, dtype=weight.dtype, out=weight)
+ else:
+ self.update(index, weight, grad, state)
+
def set_learning_rate(self, lr):
"""Sets a new learning rate of the optimizer.
@@ -399,29 +463,29 @@ class SGD(Optimizer):
in 32-bit precision even if actual weights used in the model have lower precision.\
Turning this on can improve convergence and accuracy when training with float16.
"""
- def __init__(self, momentum=0.0, multi_precision=False, **kwargs):
+ def __init__(self, momentum=0.0, **kwargs):
super(SGD, self).__init__(**kwargs)
self.momentum = momentum
- self.multi_precision = multi_precision
- def create_state(self, index, weight):
- momentum = None
+ def create_state_multi_precision(self, index, weight):
+ weight_master_copy = None
if self.multi_precision and weight.dtype == numpy.float16:
- weight_master_copy = array(weight, ctx=weight.context, dtype=numpy.float32)
- if self.momentum != 0.0:
- momentum = zeros(weight.shape, weight.context, dtype=numpy.float32,
- stype=weight.stype)
- return (momentum, weight_master_copy)
+ weight_master_copy = weight.astype(numpy.float32)
+ return (self.create_state(index, weight_master_copy), weight_master_copy)
if weight.dtype == numpy.float16 and not self.multi_precision:
warnings.warn("Accumulating with float16 in optimizer can lead to "
"poor accuracy or slow convergence. "
"Consider using multi_precision=True option of the "
"SGD optimizer")
+ return self.create_state(index, weight)
+
+ def create_state(self, index, weight):
+ momentum = None
if self.momentum != 0.0:
momentum = zeros(weight.shape, weight.context, dtype=weight.dtype, stype=weight.stype)
return momentum
- def update(self, index, weight, grad, state):
+ def _update_impl(self, index, weight, grad, state, multi_precision=False):
assert(isinstance(weight, NDArray))
assert(isinstance(grad, NDArray))
self._update_count(index)
@@ -433,9 +497,8 @@ class SGD(Optimizer):
kwargs['momentum'] = self.momentum
if self.clip_gradient:
kwargs['clip_gradient'] = self.clip_gradient
- use_multi_precision = isinstance(state, (list, tuple))
- if not use_multi_precision:
+ if not multi_precision:
if state is not None:
sgd_mom_update(weight, grad, state, out=weight,
lr=lr, wd=wd, **kwargs)
@@ -450,6 +513,14 @@ class SGD(Optimizer):
mp_sgd_update(weight, grad, state[1], out=weight,
lr=lr, wd=wd, **kwargs)
+ def update(self, index, weight, grad, state):
+ self._update_impl(index, weight, grad, state, multi_precision=False)
+
+ def update_multi_precision(self, index, weight, grad, state):
+ use_multi_precision = self.multi_precision and weight.dtype == numpy.float16
+ self._update_impl(index, weight, grad, state,
+ multi_precision=use_multi_precision)
+
# pylint: enable=line-too-long
@register
class DCASGD(Optimizer):
@@ -1029,13 +1100,13 @@ class Updater(object):
if isinstance(index, bytes):
index = py_str(index)
if index not in self.states:
- self.states[index] = self.optimizer.create_state(index, weight)
+ self.states[index] = self.optimizer.create_state_multi_precision(index, weight)
self.states_synced[index] = True
elif not self.states_synced[index]:
self.states[index] = \
self.sync_state_context(self.states[index], weight.context)
self.states_synced[index] = True
- self.optimizer.update(index, weight, grad, self.states[index])
+ self.optimizer.update_multi_precision(index, weight, grad, self.states[index])
def sync_state_context(self, state, context):
if isinstance(state, NDArray):
diff --git a/tests/python/unittest/test_optimizer.py b/tests/python/unittest/test_optimizer.py
index e55a7ec..1a56859 100644
--- a/tests/python/unittest/test_optimizer.py
+++ b/tests/python/unittest/test_optimizer.py
@@ -65,6 +65,14 @@ def test_lr_wd_mult():
assert not mx.test_utils.almost_equal(args1['fc1_bias'], args2['fc1_bias'], 1e-1)
assert not mx.test_utils.almost_equal(args1['fc2_weight'], args2['fc2_weight'], 1e-1)
+def compare_ndarray_tuple(t1, t2, rtol=None, atol=None):
+ if t1 is not None and t2 is not None:
+ if isinstance(t1, tuple):
+ for s1, s2 in zip(t1, t2):
+ compare_ndarray_tuple(s1, s2, rtol, atol)
+ else:
+ assert_almost_equal(t1.asnumpy(), t2.asnumpy(), rtol=rtol, atol=atol)
+
def compare_optimizer(opt1, opt2, shape, dtype, w_stype='default', g_stype='default'):
if w_stype == 'default':
@@ -84,25 +92,13 @@ def compare_optimizer(opt1, opt2, shape, dtype, w_stype='default', g_stype='defa
else:
raise Exception("type not supported yet")
- state1 = opt1.create_state(0, w1)
- state2 = opt2.create_state(0, w2)
- if state1 is not None and state2 is not None:
- if isinstance(state1, tuple):
- for s1, s2, in zip(state1, state2):
- if s1 is not None or s2 is not None:
- assert(same(s1.asnumpy(), s2.asnumpy()))
- else:
- assert_almost_equal(state1.asnumpy(), state2.asnumpy())
-
- opt1.update(0, w1, g1, state1)
- opt2.update(0, w2, g2, state2)
- if state1 is not None and state2 is not None:
- if isinstance(state1, tuple):
- for s1, s2, in zip(state1, state2):
- if s1 is not None or s2 is not None:
- assert_almost_equal(s1.asnumpy(), s2.asnumpy(), rtol=1e-4, atol=1e-5)
- else:
- assert_almost_equal(state1.asnumpy(), state2.asnumpy())
+ state1 = opt1.create_state_multi_precision(0, w1)
+ state2 = opt2.create_state_multi_precision(0, w2)
+ compare_ndarray_tuple(state1, state2)
+
+ opt1.update_multi_precision(0, w1, g1, state1)
+ opt2.update_multi_precision(0, w2, g2, state2)
+ compare_ndarray_tuple(state1, state2, rtol=1e-4, atol=1e-5)
assert_almost_equal(w1.asnumpy(), w2.asnumpy(), rtol=1e-4, atol=1e-5)
# SGD
@@ -136,6 +132,9 @@ class PySGD(mx.optimizer.Optimizer):
momentum = mx.nd.zeros(weight.shape, weight.context, dtype=weight.dtype)
return momentum
+ def create_state_multi_precision(self, index, weight):
+ return self.create_state(index, weight)
+
def update(self, index, weight, grad, state):
"""Update the parameters.
@@ -195,6 +194,9 @@ class PySGD(mx.optimizer.Optimizer):
tmp = weight32.astype(weight.dtype)
tmp.copyto(weight)
+ def update_multi_precision(self, index, weight, grad, state):
+ self.update(index, weight, grad, state)
+
def test_sgd():
mx.random.seed(0)
opt1 = PySGD
@@ -308,7 +310,7 @@ def test_sparse_sgd():
cg_options = [{}, {'clip_gradient': 0.4}, {'clip_gradient': 0.5}]
rg_options = [{}, {'rescale_grad': 0.14}, {'rescale_grad': 0.8}]
wd_options = [{}, {'wd': 0.03}, {'wd': 0.05}, {'wd': 0.07}]
- mp_options = [{}]
+ mp_options = [{}, {'multi_precision': False}, {'multi_precision': True}]
for dtype in [np.float32]:
for mom_option in mom_options:
for cg_option in cg_options:
@@ -402,17 +404,28 @@ def test_adam():
opt1 = PyAdam
opt2 = mx.optimizer.Adam
shape = (3, 4, 5)
- kwargs = [{},
- {'clip_gradient': 0.5},
- {'clip_gradient': 0.4, 'rescale_grad': 0.14},
- {'rescale_grad': 0.8},
- {'clip_gradient': 0.5, 'wd': 0.07},
- {'clip_gradient': 0.4, 'rescale_grad': 0.14, 'wd': 0.03},
- {'rescale_grad': 0.8, 'wd': 0.05}]
- for kwarg in kwargs:
- compare_optimizer(opt1(**kwarg), opt2(**kwarg), shape, np.float32)
- compare_optimizer(opt1(sparse_update=True, **kwarg), opt2(**kwarg), shape,
- np.float32, w_stype='row_sparse', g_stype='row_sparse')
+ cg_options = [{}, {'clip_gradient': 0.4}, {'clip_gradient': 0.5}]
+ rg_options = [{}, {'rescale_grad': 0.14}, {'rescale_grad': 0.8}]
+ wd_options = [{}, {'wd': 0.03}, {'wd': 0.05}, {'wd': 0.07}]
+ mp_options = [{}, {'multi_precision': False}, {'multi_precision': True}]
+ for dtype in [np.float16, np.float32, np.float64]:
+ for cg_option in cg_options:
+ for rg_option in rg_options:
+ for wd_option in wd_options:
+ for mp_option in mp_options:
+ kwarg = {}
+ kwarg.update(cg_option)
+ kwarg.update(rg_option)
+ kwarg.update(wd_option)
+ kwarg.update(mp_option)
+ if (dtype == np.float16 and
+ ('multi_precision' not in kwarg or
+ not kwarg['multi_precision'])):
+ continue
+ compare_optimizer(opt1(**kwarg), opt2(**kwarg), shape, dtype)
+ if (default_context() == mx.cpu()):
+ compare_optimizer(opt1(sparse_update=True, **kwarg), opt2(**kwarg), shape,
+ dtype, w_stype='row_sparse', g_stype='row_sparse')
# RMSProp
class PyRMSProp(mx.optimizer.Optimizer):
@@ -524,36 +537,33 @@ def test_rms():
opt1 = PyRMSProp
opt2 = mx.optimizer.RMSProp
shape = (3, 4, 5)
- kwargs = [{},
- {'clip_gradient': 0.5},
- {'clip_gradient': 0.4, 'rescale_grad': 0.14},
- {'rescale_grad': 0.8},
- {'clip_gradient': 0.5, 'wd': 0.07},
- {'clip_gradient': 0.4, 'rescale_grad': 0.14, 'wd': 0.03},
- {'rescale_grad': 0.8, 'wd': 0.05},
- {'centered': True},
- {'clip_gradient': 0.5, 'centered': True},
- {'clip_gradient': 0.4, 'rescale_grad': 0.14, 'centered': True},
- {'rescale_grad': 0.8, 'centered': True},
- {'clip_gradient': 0.5, 'wd': 0.07, 'centered': True},
- {'clip_gradient': 0.4, 'rescale_grad': 0.14, 'wd': 0.03, 'centered': True},
- {'rescale_grad': 0.8, 'wd': 0.05, 'centered': True},
- {'clip_gradient': 0.5, 'clip_weights': 0.01},
- {'clip_gradient': 0.4, 'rescale_grad': 0.14, 'clip_weights': 0.01},
- {'rescale_grad': 0.8, 'clip_weights': 0.01},
- {'clip_gradient': 0.5, 'wd': 0.07, 'clip_weights': 0.01},
- {'clip_gradient': 0.4, 'rescale_grad': 0.14, 'wd': 0.03, 'clip_weights': 0.01},
- {'rescale_grad': 0.8, 'wd': 0.05, 'clip_weights': 0.01},
- {'centered': True, 'clip_weights': 0.01},
- {'clip_gradient': 0.5, 'centered': True, 'clip_weights': 0.01},
- {'clip_gradient': 0.4, 'rescale_grad': 0.14, 'centered': True, 'clip_weights': 0.01},
- {'rescale_grad': 0.8, 'centered': True, 'clip_weights': 0.01},
- {'clip_gradient': 0.5, 'wd': 0.07, 'centered': True, 'clip_weights': 0.01},
- {'clip_gradient': 0.4, 'rescale_grad': 0.14, 'wd': 0.03, 'centered': True, 'clip_weights': 0.01},
- {'rescale_grad': 0.8, 'wd': 0.05, 'centered': True, 'clip_weights': 0.01}]
- for kwarg in kwargs:
- compare_optimizer(opt1(**kwarg), opt2(**kwarg), shape, np.float32)
- compare_optimizer(opt1(**kwarg), opt2(**kwarg), shape, np.float32, g_stype='row_sparse')
+ cg_options = [{}, {'clip_gradient': 0.4}, {'clip_gradient': 0.5}]
+ cw_options = [{}, {'clip_weights': 0.01}]
+ center_options = [{}, {'centered': False}, {'centered': True}]
+ rg_options = [{}, {'rescale_grad': 0.14}, {'rescale_grad': 0.8}]
+ wd_options = [{}, {'wd': 0.03}, {'wd': 0.05}, {'wd': 0.07}]
+ mp_options = [{}, {'multi_precision': False}, {'multi_precision': True}]
+ for dtype in [np.float16, np.float32]:
+ for cw_option in cw_options:
+ for cg_option in cg_options:
+ for center_option in center_options:
+ for rg_option in rg_options:
+ for wd_option in wd_options:
+ for mp_option in mp_options:
+ kwarg = {}
+ kwarg.update(cw_option)
+ kwarg.update(cg_option)
+ kwarg.update(center_option)
+ kwarg.update(rg_option)
+ kwarg.update(wd_option)
+ kwarg.update(mp_option)
+ if (dtype == np.float16 and
+ ('multi_precision' not in kwarg or
+ not kwarg['multi_precision'])):
+ continue
+ compare_optimizer(opt1(**kwarg), opt2(**kwarg), shape, dtype)
+ if (default_context() == mx.cpu()):
+ compare_optimizer(opt1(**kwarg), opt2(**kwarg), shape, dtype, g_stype='row_sparse')
class PyFtrl(mx.optimizer.Optimizer):
"""The Ftrl optimizer.
--
To stop receiving notification emails like this one, please contact
['"commits@mxnet.apache.org" <co...@mxnet.apache.org>'].