Source code for mxnet.optimizer.nadam

# coding: utf-8
# 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.

# pylint: disable=W0223
"""Nadam optimizer."""
from __future__ import absolute_import
from ..ndarray import (zeros, clip, sqrt, square)
from .optimizer import Optimizer, register

__all__ = ['Nadam']


[docs]@register class Nadam(Optimizer): """The Nesterov Adam optimizer. Much like Adam is essentially RMSprop with momentum, Nadam is Adam RMSprop with Nesterov momentum available at http://cs229.stanford.edu/proj2015/054_report.pdf. This optimizer accepts the following parameters in addition to those accepted by :class:`.Optimizer`. Parameters ---------- learning_rate : float, default 0.001 The initial learning rate. If None, the optimization will use the learning rate from ``lr_scheduler``. If not None, it will overwrite the learning rate in ``lr_scheduler``. If None and ``lr_scheduler`` is also None, then it will be set to 0.01 by default. beta1 : float, default 0.9 Exponential decay rate for the first moment estimates. beta2 : float, default 0.999 Exponential decay rate for the second moment estimates. epsilon : float, default 1e-8 Small value to avoid division by 0. schedule_decay : float, default 0.004 Exponential decay rate for the momentum schedule use_fused_step : bool, default False Whether or not to use fused kernels for optimizer. When use_fused_step=False, step is called, otherwise, fused_step is called. """ def __init__(self, learning_rate=0.001, beta1=0.9, beta2=0.999, epsilon=1e-8, schedule_decay=0.004, use_fused_step=False, **kwargs): super(Nadam, self).__init__(learning_rate=learning_rate, use_fused_step=use_fused_step, **kwargs) self.beta1 = beta1 self.beta2 = beta2 self.epsilon = epsilon self.schedule_decay = schedule_decay self.m_schedule = 1.
[docs] def create_state(self, index, weight): return (zeros(weight.shape, weight.context, dtype=weight.dtype), # mean zeros(weight.shape, weight.context, dtype=weight.dtype)) # variance
[docs] def step(self, indices, weights, grads, states): """Perform an optimization step using gradients and states. Parameters ---------- indices : list of int List of unique indices of the parameters 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. weights : list of NDArray List of parameters to be updated. grads : list of NDArray List of gradients of the objective with respect to this parameter. states : List of any obj List of state returned by `create_state()`. """ for index, weight, grad, state in zip(indices, weights, grads, states): self._update_count(index) lr = self._get_lr(index) wd = self._get_wd(index) t = self._index_update_count[index] # preprocess grad grad *= self.rescale_grad if self.clip_gradient is not None: grad = clip(grad, -self.clip_gradient, self.clip_gradient) grad += wd * weight coef2 = 1. - self.beta2**t # warming momentum schedule momentum_t = self.beta1 * (1. - 0.5 * (pow(0.96, t * self.schedule_decay))) momentum_t_1 = self.beta1 * (1. - 0.5 * (pow(0.96, (t + 1) * self.schedule_decay))) self.m_schedule = self.m_schedule * momentum_t m_schedule_next = self.m_schedule * momentum_t_1 # update mean and var mean, var = state mean[:] *= self.beta1 mean[:] += (1. - self.beta1) * grad var[:] *= self.beta2 var[:] += (1. - self.beta2) * square(grad) grad_prime = grad / (1. - self.m_schedule) mean_prime = mean / (1. - m_schedule_next) var_prime = var / coef2 mean_bar = momentum_t_1 * mean_prime + (1. - momentum_t) * grad_prime # update weight d = mean_bar / (sqrt(var_prime) + self.epsilon) weight[:] -= lr * d