| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, |
| # software distributed under the License is distributed on an |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| # KIND, either express or implied. See the License for the |
| # specific language governing permissions and limitations |
| # under the License. |
| |
| # coding: utf-8 |
| # pylint: disable=line-too-long |
| """Parameter optimizer.""" |
| __all__ = ['Trainer'] |
| |
| import warnings |
| from collections import OrderedDict |
| |
| from .. import optimizer as opt |
| from ..model import _create_kvstore, _create_sparse_kvstore |
| from .parameter import Parameter |
| from ..kvstore import KVStore |
| |
| |
| class Trainer(object): |
| """Applies an `Optimizer` on a set of Parameters. Trainer should |
| be used together with `autograd`. |
| |
| .. note:: |
| |
| For the following cases, updates will always happen on kvstore, |
| i.e., you cannot set update_on_kvstore=False. |
| |
| - dist kvstore with sparse weights or sparse gradients |
| - dist async kvstore |
| - `optimizer.lr_scheduler` is not None |
| |
| Parameters |
| ---------- |
| params : Dict |
| The set of parameters to optimize. |
| optimizer : str or Optimizer |
| The optimizer to use. See |
| `help <https://mxnet.apache.org/api/python/docs/api/optimizer/index.html>`_ |
| on Optimizer for a list of available optimizers. |
| optimizer_params : dict |
| Key-word arguments to be passed to optimizer constructor. For example, |
| `{'learning_rate': 0.1}`. All optimizers accept learning_rate, wd (weight decay), |
| clip_gradient, and lr_scheduler. See each optimizer's |
| constructor for a list of additional supported arguments. |
| kvstore : str or KVStore |
| kvstore type for multi-gpu and distributed training. See help on |
| :func:`mxnet.kvstore.create` for more information. |
| compression_params : dict |
| Specifies type of gradient compression and additional arguments depending |
| on the type of compression being used. For example, 2bit compression requires a threshold. |
| Arguments would then be {'type':'2bit', 'threshold':0.5} |
| See mxnet.KVStore.set_gradient_compression method for more details on gradient compression. |
| update_on_kvstore : bool, default None |
| Whether to perform parameter updates on kvstore. If None and optimizer.aggregate_num <= 1, |
| then trainer will choose the more suitable option depending on the type of kvstore. |
| If None and optimizer.aggregate_num > 1, `update_on_kvstore` is set to False. |
| If the `update_on_kvstore` argument is provided, |
| environment variable `MXNET_UPDATE_ON_KVSTORE` will be ignored. |
| |
| Properties |
| ---------- |
| learning_rate : float |
| The current learning rate of the optimizer. Given an Optimizer object |
| optimizer, its learning rate can be accessed as optimizer.learning_rate. |
| """ |
| def __init__(self, params, optimizer, optimizer_params=None, kvstore='device', |
| compression_params=None, update_on_kvstore=None): |
| param_list = [] |
| if isinstance(params, (dict, OrderedDict)): |
| for key in sorted(list(params.keys())): |
| param_list.append(params[key]) |
| params = param_list |
| if not isinstance(params, (list, tuple)): |
| raise ValueError( |
| "First argument must be a list or dict of Parameters, " \ |
| f"got {type(params)}.") |
| self._params = [] |
| # parameters to initialize on the kvstore |
| self._contains_sparse_weight = False |
| self._contains_sparse_grad = False |
| self._param2idx = {} |
| for i, param in enumerate(params): |
| if not isinstance(param, Parameter): |
| raise ValueError( |
| "First argument must be a list or dict of Parameters, " \ |
| f"got list of {type(param)}.") |
| if param._uuid in self._param2idx: |
| # Shared parameters have same uuid; only need to store one of the shared versions |
| continue |
| self._param2idx[param._uuid] = i |
| self._params.append(param) |
| param._set_trainer(self) |
| if param._stype != 'default': |
| self._contains_sparse_weight = True |
| if param._grad_stype != 'default': |
| self._contains_sparse_grad = True |
| self._compression_params = compression_params |
| self._devices = self._check_devices() |
| optimizer_params = optimizer_params if optimizer_params else {} |
| self._init_optimizer(optimizer, optimizer_params) |
| self._scale = self._optimizer.rescale_grad |
| if self._optimizer.aggregate_num > 1 and update_on_kvstore is not None: |
| if update_on_kvstore: |
| raise ValueError("Cannot set update_on_kvstore=True " |
| "when optimizer.aggregate_num > 1.") |
| if update_on_kvstore is None and self._optimizer.aggregate_num > 1: |
| update_on_kvstore = False |
| self._kvstore_params = {'kvstore': kvstore, 'update_on_kvstore': update_on_kvstore} |
| self._kv_initialized = False |
| self._kvstore = None |
| self._update_on_kvstore = None |
| self._distributed = None |
| self._params_to_init = [] |
| self._reset_kvstore() |
| |
| def _check_contexts(self): |
| """This function has been deprecated. Please refer to ``Trainer._check_devices``.""" |
| warnings.warn('Trainer._check_contexts has been renamed to' |
| ' Trainer._check_devices', DeprecationWarning) |
| return self._check_devices() |
| |
| def _check_devices(self): |
| devices = None |
| for param in self._params: |
| device = param.list_device() |
| assert devices is None or devices == device, \ |
| "All Parameters must be initialized on the same set of devices, " \ |
| f"but Parameter {param.name} is initialized on {str(device)} while previous Parameters " \ |
| f"are initialized on {str(devices)}." |
| devices = device |
| return devices |
| |
| def _init_optimizer(self, optimizer, optimizer_params): |
| param_dict = {i: param for i, param in enumerate(self._params)} |
| if isinstance(optimizer, opt.Optimizer): |
| assert not optimizer_params, \ |
| "optimizer_params must be None if optimizer is an instance of " \ |
| "Optimizer instead of str" |
| self._optimizer = optimizer |
| # param_dict must not be deep copied, so that if user mutate the lr_mult |
| # or wd_mult of some parameters, it takes effect. |
| self._optimizer.param_dict = param_dict |
| else: |
| self._optimizer = opt.create(optimizer, param_dict=param_dict, |
| **optimizer_params) |
| self._updaters = [opt.get_updater(self._optimizer) \ |
| for _ in self._devices] |
| |
| def _init_params(self): |
| """Initialize parameters in the KVStore. |
| |
| Parameters with incomplete initialization are ignored. |
| |
| """ |
| assert self._kv_initialized, "Cannot initialize parameters in KVStore " \ |
| "when KVStore is not initialized." |
| params_to_init = [] |
| if self._kvstore: |
| for param in self._params_to_init: |
| if param._deferred_init: |
| params_to_init.append(param) |
| else: |
| param_arrays = param._check_and_get(param._data, list) |
| idx = self._param2idx[param._uuid] |
| if param._stype != 'default': |
| self._kvstore.init(idx, param_arrays[0]) |
| else: |
| self._kvstore.broadcast(idx, param_arrays[0], param_arrays) |
| |
| self._params_to_init = params_to_init |
| |
| def _reset_kvstore(self): |
| """Reset kvstore.""" |
| if self._kvstore and 'dist' in self._kvstore.type: |
| raise RuntimeError("Cannot reset distributed KVStore.") |
| self._kv_initialized = False |
| self._kvstore = None |
| self._distributed = None |
| self._update_on_kvstore = None |
| self._params_to_init = [param for param in self._params] |
| |
| def _init_kvstore(self): |
| """Create kvstore.""" |
| config = self._kvstore_params |
| # configure kvstore, update_on_kvstore and self._distributed on three cases: |
| if self._contains_sparse_weight: |
| # If weight is sparse, kvstore must be present and the weight must be updated on kvstore. |
| # The training loop is the following: |
| # - row_sparse_pull(sparse_weight) |
| # - forward() |
| # - backward() |
| # - push_and_update(grad) |
| # - pull(weight) |
| kvstore, update_on_kvstore = _create_sparse_kvstore(config['kvstore']) |
| self._distributed = 'dist' in kvstore.type |
| # raise err if user provides unsupported configs |
| if config['update_on_kvstore'] is False: |
| raise ValueError("Cannot set update_on_kvstore=False when sparse weights " |
| "are present.") |
| |
| elif self._contains_sparse_grad: |
| # For single node training with dense weight and sparse grad, |
| # we prefer update_on_kvstore=False because this is usually faster. |
| # This means we push and pull sparse gradients, and we do not store weight in kvstore. |
| # The training loop is the following: |
| # - forward() |
| # - backward() |
| # - push(grad) |
| # - pull(grad) |
| # - update(grad, weight) |
| # |
| # For multi-node training with dense weight and sparse grad, |
| # only update_on_kvstore=True is supported, due to the fact that |
| # kv.row_sparse_pull(grad) is not implemented. |
| # Therefore, we push sparse gradients and pull dense weights. |
| # The training loop contains: |
| # - forward() |
| # - backward() |
| # - push_and_update(grad) |
| # - pull(weight) |
| arg_arrays = {param._uuid: param.data(self._devices[0]) for param in self._params} |
| kvstore, _ = _create_kvstore(config['kvstore'], len(self._devices), arg_arrays) |
| self._distributed = 'dist' in kvstore.type if kvstore else False |
| update_on_kvstore = self._distributed |
| # raise err if user provides unsupported configs |
| if config['update_on_kvstore'] is not None: |
| if config['update_on_kvstore'] is False and self._distributed: |
| raise ValueError("Cannot set update_on_kvstore=False on dist kvstore " |
| "when sparse gradients are present.") |
| update_on_kvstore = config['update_on_kvstore'] |
| # raise err if a custom kvstore is used for sparse training |
| if kvstore is not None and not isinstance(kvstore, KVStore): |
| raise ValueError("Cannot use {} for multi-device training with sparse gradients" |
| .format(type(kvstore))) |
| |
| else: |
| # Training with dense weight and dense gradients. |
| # The only unsupported mode is async with update_on_kvstore=False |
| arg_arrays = {param._uuid: param.data(self._devices[0]) for param in self._params} |
| kvstore, update_on_kvstore = _create_kvstore(config['kvstore'], len(self._devices), |
| arg_arrays) |
| self._distributed = 'dist' in kvstore.type if kvstore else False |
| if self._distributed and 'async' in kvstore.type: |
| update_on_kvstore = True |
| # raise err if user provides unsupported configs |
| if config['update_on_kvstore'] is False: |
| raise ValueError("Please set update_on_kvstore=True " |
| "when training in async mode.") |
| if config['update_on_kvstore'] is not None: |
| update_on_kvstore = config['update_on_kvstore'] |
| # raise err if update_on_kvstore is set to True with kvstores that do not support optimizers |
| if update_on_kvstore and not kvstore.is_capable('optimizer'): |
| if config['update_on_kvstore']: |
| raise ValueError("Please set update_on_kvstore=False " |
| "when training with {}".format(type(kvstore))) |
| update_on_kvstore = False |
| |
| # set grad compression and optimizers |
| if kvstore: |
| if self._compression_params: |
| kvstore.set_gradient_compression(self._compression_params) |
| if update_on_kvstore: |
| # optimizer preferably needs to be set before init for multiprecision |
| kvstore.set_optimizer(self._optimizer) |
| self._kvstore = kvstore |
| self._update_on_kvstore = update_on_kvstore |
| else: |
| self._kvstore = None |
| self._update_on_kvstore = None |
| |
| self._kv_initialized = True |
| |
| @property |
| def learning_rate(self): |
| if not isinstance(self._optimizer, opt.Optimizer): |
| raise UserWarning("Optimizer has to be defined before its learning " |
| "rate can be accessed.") |
| |
| return self._optimizer.learning_rate |
| |
| @property |
| def optimizer(self): |
| if isinstance(self._optimizer, opt.Optimizer): |
| return self._optimizer |
| else: |
| raise UserWarning("Optimizer has not been initialized yet") |
| |
| def set_learning_rate(self, lr): |
| """Sets a new learning rate of the optimizer. |
| |
| Parameters |
| ---------- |
| lr : float |
| The new learning rate of the optimizer. |
| """ |
| if not isinstance(self._optimizer, opt.Optimizer): |
| raise UserWarning("Optimizer has to be defined before its learning " |
| "rate is mutated.") |
| |
| self._optimizer.set_learning_rate(lr) |
| |
| def _row_sparse_pull(self, parameter, out, row_id, full_idx=False): |
| """Internal method to invoke pull operations on KVStore. If `full_idx` is set to True, |
| `kv.pull` is preferred instead of `kv.row_sparse_pull`. |
| """ |
| # initialize kv and params if not already |
| if not self._kv_initialized: |
| self._init_kvstore() |
| if self._params_to_init: |
| self._init_params() |
| idx = self._param2idx[parameter._uuid] |
| if full_idx and 'dist' not in self._kvstore.type: |
| assert row_id.size == out.shape[0] |
| self._kvstore.pull(idx, out=out, priority=-idx, ignore_sparse=False) |
| else: |
| self._kvstore.row_sparse_pull(idx, out=out, row_ids=row_id, priority=-idx) |
| |
| def _check_and_rescale_grad(self, scale): |
| if self._update_on_kvstore and self._distributed and self._kv_initialized: |
| if self._optimizer.rescale_grad != scale: |
| raise UserWarning('Possible change in the `batch_size` from previous ' |
| '`step` detected. Optimizer gradient normalizing ' |
| 'factor will not change w.r.t new batch_size when ' |
| 'update_on_kvstore=True and when distributed kvstore ' |
| 'is used.') |
| self._optimizer.rescale_grad = scale |
| |
| def step(self, batch_size, ignore_stale_grad=False): |
| """Makes one step of parameter update. Should be called after |
| `autograd.backward()` and outside of `record()` scope. |
| |
| For normal parameter updates, `step()` should be used, which internally calls |
| `allreduce_grads()` and then `update()`. However, if you need to get the reduced |
| gradients to perform certain transformation, such as in gradient clipping, then |
| you may want to manually call `allreduce_grads()` and `update()` separately. |
| |
| Parameters |
| ---------- |
| batch_size : int |
| Batch size of data processed. Gradient will be normalized by `1/batch_size`. |
| Set this to 1 if you normalized loss manually with `loss = mean(loss)`. |
| ignore_stale_grad : bool, optional, default=False |
| If true, ignores Parameters with stale gradient (gradient that has not |
| been updated by `backward` after last step) and skip update. |
| """ |
| rescale_grad = self._scale / batch_size |
| self._check_and_rescale_grad(rescale_grad) |
| |
| if not self._kv_initialized: |
| self._init_kvstore() |
| if self._params_to_init: |
| self._init_params() |
| |
| self._allreduce_grads() |
| self._update(ignore_stale_grad) |
| |
| def allreduce_grads(self): |
| """For each parameter, reduce the gradients from different devices. |
| |
| Should be called after `autograd.backward()`, outside of `record()` scope, |
| and before `trainer.update()`. |
| |
| For normal parameter updates, `step()` should be used, which internally calls |
| `allreduce_grads()` and then `update()`. However, if you need to get the reduced |
| gradients to perform certain transformation, such as in gradient clipping, then |
| you may want to manually call `allreduce_grads()` and `update()` separately. |
| """ |
| if not self._kv_initialized: |
| self._init_kvstore() |
| if self._params_to_init: |
| self._init_params() |
| assert not (self._kvstore and self._update_on_kvstore), \ |
| 'allreduce_grads() when parameters are updated on kvstore ' \ |
| 'is not supported. Try setting `update_on_kvstore` ' \ |
| 'to False when creating trainer.' |
| |
| self._allreduce_grads() |
| |
| def _allreduce_grads(self): |
| # nothing to reduce |
| if not self._kvstore: |
| return |
| for i, param in enumerate(self._params): |
| if param.grad_req != 'null': |
| idx = self._param2idx[param._uuid] |
| grad_list = param.list_grad() |
| # sparse gradients, call push and pull separately |
| if grad_list[0].stype != 'default': |
| self._kvstore.push(idx, grad_list, priority=-i) |
| if param._stype == 'default': |
| if self._update_on_kvstore: |
| pull_list = param.list_data() |
| else: |
| pull_list = param.list_grad() |
| self._kvstore.pull(idx, pull_list, priority=-i, |
| ignore_sparse=self._distributed) |
| else: |
| # allreduce dense gradients if not update_on_kvstore, |
| # otherwise push dense gradients, pull dense weights |
| if self._update_on_kvstore: |
| self._kvstore.pushpull(idx, grad_list, out=param.list_data(), priority=-i) |
| else: |
| self._kvstore.pushpull(idx, grad_list, priority=-i) |
| |
| def update(self, batch_size, ignore_stale_grad=False): |
| """Makes one step of parameter update. |
| |
| Should be called after `autograd.backward()` and outside of `record()` scope, |
| and after `trainer.update()`. |
| |
| |
| For normal parameter updates, `step()` should be used, which internally calls |
| `allreduce_grads()` and then `update()`. However, if you need to get the reduced |
| gradients to perform certain transformation, such as in gradient clipping, then |
| you may want to manually call `allreduce_grads()` and `update()` separately. |
| |
| Parameters |
| ---------- |
| batch_size : int |
| Batch size of data processed. Gradient will be normalized by `1/batch_size`. |
| Set this to 1 if you normalized loss manually with `loss = mean(loss)`. |
| ignore_stale_grad : bool, optional, default=False |
| If true, ignores Parameters with stale gradient (gradient that has not |
| been updated by `backward` after last step) and skip update. |
| """ |
| if not self._kv_initialized: |
| self._init_kvstore() |
| if self._params_to_init: |
| self._init_params() |
| assert not (self._kvstore and self._update_on_kvstore), \ |
| 'update() when parameters are updated on kvstore ' \ |
| 'is not supported. Try setting `update_on_kvstore` ' \ |
| 'to False when creating trainer.' |
| |
| self._check_and_rescale_grad(self._scale / batch_size) |
| self._update(ignore_stale_grad) |
| |
| def _update(self, ignore_stale_grad=False): |
| loss_scaler = getattr(self, '_amp_loss_scaler', None) |
| if loss_scaler is not None: |
| if loss_scaler.has_overflow(self._params): |
| return # skip on overflow |
| |
| updates = [[] for _ in self._updaters] |
| |
| for i, param in enumerate(self._params): |
| if param.grad_req == 'null': |
| continue |
| |
| if not ignore_stale_grad: |
| for data in param._check_and_get(param._data, list): |
| if not data._fresh_grad: |
| raise UserWarning( |
| f"Gradient of Parameter `{param.name}` on device {str(data.device)} has not been updated " |
| "by backward since last `step`. This could mean a bug in your " |
| "model that made it only use a subset of the Parameters (Blocks) " |
| "for this iteration. If you are intentionally only using a subset, " |
| "call step with ignore_stale_grad=True to suppress this " |
| "warning and skip updating of Parameters with stale gradient") |
| |
| if self._kvstore and self._update_on_kvstore: |
| continue |
| |
| for upd, arr, grad in zip(updates, param.list_data(), param.list_grad()): |
| if not ignore_stale_grad or arr._fresh_grad: |
| upd.append((i, grad, arr)) |
| arr._fresh_grad = False |
| |
| if not (self._kvstore and self._update_on_kvstore): |
| for updater, upd in zip(self._updaters, updates): |
| if upd: |
| i, g, w = zip(*upd) |
| updater(i, g, w) |
| |
| def save_states(self, fname): |
| """Saves trainer states (e.g. optimizer, momentum) to a file. |
| |
| |
| Parameters |
| ---------- |
| fname : str |
| Path to output states file. |
| |
| Note |
| ---- |
| `optimizer.param_dict`, which contains Parameter information (such as |
| `lr_mult` and `wd_mult`) will not be saved. |
| """ |
| assert self._optimizer is not None |
| |
| if not self._kv_initialized: |
| self._init_kvstore() |
| if self._params_to_init: |
| self._init_params() |
| |
| if self._update_on_kvstore: |
| assert not self._params_to_init, "Cannot save trainer states when some " \ |
| "parameters are not yet initialized in kvstore." |
| self._kvstore.save_optimizer_states(fname, dump_optimizer=True) |
| else: |
| with open(fname, 'wb') as fout: |
| fout.write(self._updaters[0].get_states(dump_optimizer=True)) |
| |
| def load_states(self, fname): |
| """Loads trainer states (e.g. optimizer, momentum) from a file. |
| |
| Parameters |
| ---------- |
| fname : str |
| Path to input states file. |
| |
| Note |
| ---- |
| `optimizer.param_dict`, which contains Parameter information (such as |
| `lr_mult` and `wd_mult`) will not be loaded from the file, but rather set |
| based on current Trainer's parameters. |
| """ |
| if not self._kv_initialized: |
| self._init_kvstore() |
| if self._params_to_init: |
| self._init_params() |
| |
| if self._update_on_kvstore: |
| self._kvstore.load_optimizer_states(fname) |
| self._optimizer = self._kvstore._updater.optimizer |
| else: |
| with open(fname, 'rb') as f: |
| states = f.read() |
| for updater in self._updaters: |
| updater.set_states(states) |
| updater.optimizer = self._updaters[0].optimizer |
| self._optimizer = self._updaters[0].optimizer |
| param_dict = {i: param for i, param in enumerate(self._params)} |
| self._optimizer.param_dict = param_dict |
