mojolearn.NAdam#

class mojolearn.NAdam(params, lr=0.002, betas=(0.9, 0.999), eps=1e-08, weight_decay=0.0, momentum_decay=0.004, decoupled_weight_decay=False, numeric_mode=None)[source]#

Bases: _SeqOptimizer

NAdam optimizer, as torch.optim.NAdam.

Adam with Nesterov momentum and torch’s momentum schedule mu_t = beta1 * (1 - 0.5 * 0.96**(t * momentum_decay)); the denominator is sqrt(v / (1 - beta2**t)) + eps. Weight decay is coupled by default, decoupled (AdamW style) with decoupled_weight_decay=True. Defaults are torch’s.

Parameters:

  • params (ndarray or list of ndarray) – Parameters to update in place.

  • lr (float, default 2e-3) – Learning rate.

  • betas (tuple of float, default (0.9, 0.999)) – Decay rates of the first and second moments.

  • eps (float, default 1e-8) – Term added to the denominator.

  • weight_decay (float, default 0.0) – Weight decay.

  • momentum_decay (float, default 4e-3) – Momentum schedule decay. Must be >= 0, otherwise ValueError.

  • decoupled_weight_decay (bool, default False) – Apply weight decay to the parameters directly instead of the gradient.

  • numeric_mode ({‘identical’, ‘deterministic’, ‘fast’} or None, default None) – Numeric mode of the update. None uses the library default.

Variables:

t (int) – Number of steps taken.

Notes

params must be C-contiguous float32 NumPy arrays, updated in place; anything else raises TypeError. step(grads) takes one gradient per parameter, in the same order and shapes; float64 gradients raise TypeError. lr is read at every step, so you can change opt.lr between steps, or set opt.lr_schedule to a schedule (StepLR, ExponentialLR, OneCycleLR, or any object with a one-based lr_at(t)), which sets lr before each step. state_dict() and load_state_dict() save and restore the step count, learning rate and state buffers.

Not supported: parameter groups, maximize, foreach/fused/capturable/ differentiable, sparse gradients and float64 parameters.

Examples

>>> import numpy as np, mojolearn
>>> w = np.zeros(128, dtype=np.float32)
>>> opt = mojolearn.NAdam([w], lr=2e-3).step([np.ones_like(w)])
__init__(params, lr=0.002, betas=(0.9, 0.999), eps=1e-08, weight_decay=0.0, momentum_decay=0.004, decoupled_weight_decay=False, numeric_mode=None)[source]#

Methods

load_state_dict(sd)

state_dict()

step(grads)

One update of every param from grads (same shapes, same order).

Attributes

resident_

True once the state lives on the device.

state

On this page