pytensor_ml.optim.adam

pytensor_ml.optim.adam(learning_rate=0.001, beta1=0.9, beta2=0.999, epsilon=1e-08, amsgrad=False, *, namespace='adam')

Adam optimizer. See adam_updates() for the update rule.

learning_rate accepts a float, a scalar shared variable, any scalar graph, or a schedule, and namespace prefixes the state this rule allocates; see sgd().

Examples

The usual first choice: a per-parameter rate adapted from the first and second gradient moments, both bias-corrected, so the earliest steps are not damped towards zero:

import numpy as np

from pytensor_ml.layers import Input, Linear
from pytensor_ml.loss import SquaredError, supervised_loss
from pytensor_ml.optim import adam, compile_train

X = Input("X", shape=(None, 4))
loss, target = supervised_loss(Linear("fc", n_in=4, n_out=1)(X), SquaredError())

step = compile_train(loss, adam(learning_rate=1e-3))
loss_value = step(np.zeros((8, 4)), np.zeros((8, 1)))