mxnet.optimizer.adamW

AdamW optimizer.

Classes

AdamW([learning_rate, beta1, beta2, ...])

The AdamW optimizer.
class mxnet.optimizer.adamW.AdamW(learning_rate=0.001, beta1=0.9, beta2=0.999, epsilon=1e-06, correct_bias=True, use_fused_step=True, **kwargs)[source]

Bases: Optimizer

The AdamW optimizer.

This class implements the optimizer described in Decoupled Weight Decay Regularization,

available at https://arxiv.org/pdf/1711.05101.pdf.

Updates are applied by:

grad = clip(grad * rescale_grad, clip_gradient)
m = beta1 * m + (1 - beta1) * grad
v = beta2 * v + (1 - beta2) * (grad**2)
lr = learning_rate * sqrt(1 - beta2**t) / (1 - beta1**t)
w = w - lr * (m / (sqrt(v) + epsilon) + wd * w)

Also, we can turn off the bias correction term and the updates are as follows:

grad = clip(grad * rescale_grad, clip_gradient) + wd * weight
m = beta1 * m + (1 - beta1) * grad
v = beta2 * v + (1 - beta2) * (grad**2)
lr = learning_rate
w = w - lr * (m / (sqrt(v) + epsilon) + wd * w)

This optimizer accepts the following parameters in addition to those accepted by 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-6) – Small value to avoid division by 0.

  • correct_bias (bool, default True) – Can be set to False to avoid correcting bias in Adam (e.g. like in Bert TF repository). Default True.

  • use_fused_step (bool, default True) – Whether or not to use fused kernels for optimizer. When use_fused_step=False, step is called, otherwise, fused_step is called.

create_state(index, weight)[source]

state creation function.

fused_step(indices, weights, grads, states)[source]

Perform a fused optimization step using gradients and states. Fused kernel is used for update.

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().

step(indices, weights, grads, states)[source]

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().