mxnet.optimizer.lars

LARS optimizer.

Classes

LARS([learning_rate, momentum, eta, ...])

the LARS optimizer from 'Large Batch Training of Convolution Networks' (https://arxiv.org/abs/1708.03888)

class mxnet.optimizer.lars.LARS(learning_rate=0.1, momentum=0.0, eta=0.001, epsilon=1e-08, lazy_update=False, use_fused_step=True, aggregate_num=1, **kwargs)[source]

Bases: Optimizer

the LARS optimizer from ‘Large Batch Training of Convolution Networks’ (https://arxiv.org/abs/1708.03888)

Behave mostly like SGD with momentum and weight decay but is scaling adaptively the learning for each layer:

w_norm = L2norm(weights)
g_norm = L2norm(gradients)
if w_norm > 0 and g_norm > 0:
    lr_layer = lr * w_norm / (g_norm + weight_decay * w_norm + epsilon)
else:
    lr_layer = lr
Parameters:
  • learning_rate (float, default 0.1) – 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.

  • momentum (float, default 0.) – The momentum value.

  • eta (float, default 0.001) – LARS coefficient used to scale the learning rate.

  • epsilon (float, default 1e-8) – Small value to avoid division by 0.

  • lazy_update (bool, default False) – Default is False. If True, lazy updates are applied if the storage types of weight and grad are both row_sparse.

  • aggregate_num (int, default 1) – Number of weights to be aggregated in a list. They are passed to the optimizer for a single optimization step.

  • 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]

Creates auxiliary state for a given weight.

Some optimizers require additional states, e.g. as momentum, in addition to gradients in order to update weights. This function creates state for a given weight which will be used in update. This function is called only once for each weight.

Parameters:
  • index (int) – An unique index to identify the weight.

  • weight (NDArray) – The weight.

Returns:

state – The state associated with the weight.

Return type:

any obj

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

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

Override update_multi_precision.