mojolearn.Adamax

class mojolearn.Adamax(params, lr=0.002, betas=(0.9, 0.999), eps=1e-08, weight_decay=0.0, numeric_mode=None)

Bases: _SeqOptimizer

Adamax optimizer, as torch.optim.Adamax.

The infinity-norm variant of Adam: m = lerp(m, g, 1 - beta1), u = max(beta2 * u, |g| + eps), p -= lr / (1 - beta1**t) * m / u. Weight decay is coupled (added to the gradient). 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 moment and the infinity norm, each in [0, 1).

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

• weight_decay (float, default 0.0) – L2 penalty added to 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.Adamax([w]).step([np.ones_like(w)])

__init__(params, lr=0.002, betas=(0.9, 0.999), eps=1e-08, weight_decay=0.0, numeric_mode=None)

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