mojolearn.Adam#

class mojolearn.Adam(params, lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0.0, lr_schedule=None, accumulation_steps=1, maximize=False, resident=None, **kwargs)[source]#

Bases: _Optimizer

Adam optimizer, as torch.optim.Adam, over float32 buffers (NumPy arrays, array.array or mojolearn.Array; any writable buffer-protocol object).

Weight decay is coupled, as in torch’s Adam: it is added to the gradient and so passes through both moments. AdamW is the decoupled form; at weight_decay=0 the two compute the same thing. The bias corrections are computed in float32, not in torch’s float64, so results are close to torch’s but not bit-equal to them. Parameters are updated in place.

Parameters:

  • params (buffer or list of buffers) – float32 parameter arrays, updated in place.

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

  • betas (tuple of float, default (0.9, 0.999)) – Decay rates of the first and second moments, each in [0, 1); 1.0 is refused.

  • eps (float, default 1e-8) – Added to sqrt(v) in the denominator (sqrt(v) + eps).

  • weight_decay (float, default 0.0) – L2 penalty, added to the gradient.

  • lr_schedule (object or None, default None) – A schedule with a one-based lr_at(t) method; it sets lr before every step. Anything without lr_at raises TypeError.

  • accumulation_steps (int, default 1) – Number of microbatches combined by step_accumulated; a power of two >= 1, otherwise ValueError.

  • maximize (bool, default False) – Maximize the objective (the gradient’s sign is flipped inside the step; the caller’s gradient arrays are not negated). Must be a bool.

  • resident (bool or None, default None) – Keep the optimizer state on the device between steps when the install supports it (None or True). False keeps it on the host, which suits an optimizer that is rebuilt every step.

Variables:

  • t (int) – Number of steps taken.

  • exp_avg (array) – First moment, flat over all parameters.

  • exp_avg_sq (array) – Second moment, flat over all parameters.

Notes

step(grads, max_norm=None) takes one gradient per parameter, matching params in count, shape and order. With max_norm set, the global L2 gradient norm is clipped first: the caller’s gradient arrays are scaled in place (as torch.nn.utils.clip_grad_norm_) and step returns the pre-clip total norm; without it, step returns None and the gradients are untouched. With clipping on, the order of params is part of the result, because it fixes the summation order of the norm.

lr is read at every step, so you can assign opt.lr between steps, or pass an lr_schedule (ConstantLR, WarmupLinearLR, WarmupCosineLR, or any object with a one-based lr_at(t)). step_accumulated(microbatch_grads, tokens) combines accumulation_steps microbatch gradient lists before one step. zero_grad(grads) writes +0.0 into the given gradient arrays. state_dict() / load_state_dict() save and restore the step count and the optimizer state.

Refused with TypeError: foreach, fused, capturable, differentiable, parameter groups, sparse gradients, float16/bfloat16 and float64 buffers, and any unknown keyword. amsgrad is not supported and raises TypeError.

Examples

>>> import numpy as np, mojolearn
>>> w = np.zeros((64, 8), dtype=np.float32)
>>> opt = mojolearn.Adam([w], lr=1e-3)
>>> opt.step([np.ones_like(w)])
__init__(params, lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0.0, lr_schedule=None, accumulation_steps=1, maximize=False, resident=None, **kwargs)[source]#

Methods

get_params([deep])

grads_to_device(grads)

Upload a gradient into the handle's device gradient buffer, for step_device().

grads_to_host(grads)

Download the device gradient (clipped, if the last step clipped) into grads, in place.

load_state_dict(state)

Restore the optimizer state saved by state_dict.

numeric_mode_used()

The numeric mode this estimator runs in.

param_device(j)

Parameter tensor j as a TrainDeviceArray (a device copy of the resident parameters; params_to_device() first), for the next forward.

params_resident_available()

Whether params_to_device works on this build.

params_to_device()

Upload the parameters into the optimizer's device handle and keep them there: later steps neither upload nor download them.

params_to_host()

Download the device parameters into the caller's arrays (in place); they stay resident.

set_params(**params)

Set constructor parameters.

state_dict()

Everything a resumed run needs, and nothing that is not state.

step(grads[, max_norm])

One step.

step_accumulated(microbatch_grads, tokens[, ...])

One step from A = self.accumulation_steps microbatch gradient lists, combined on device by the clause 9.2 balanced tree in ascending microbatch index, then step.

step_device([grads, max_norm])

step on device-resident parameters (params_to_device() first).

vendor_used()

The GPU backend this estimator runs on.

zero_grad(grads)

torch.optim.Optimizer.zero_grad, over arrays the caller owns.

Attributes

exp_avg

Adam's first moment (SGD's momentum buffer), a mojolearn.Array of n_total floats.

exp_avg_sq

Adam's second moment; see exp_avg.

resident_

True once the moments live on the device.

buf_initialized

SGD's per-tensor buf_initialized flag, contract 7.3b.

t

ONE-BASED.

total_norm_

The pre-clip global gradient norm of the most recent step, or None when that step ran with max_norm=None.

clip_coef_

The clamped coefficient the most recent step applied, or None.
On this page