optimizer that implements the adam algorithm adam optimization is a stochastic gradient descent method that is based on adaptive estimation of firstorder and secondorder moments according to kingma et al 2014http arxiv orgabs1412 6980 the method is computationally efficient has little memory requirement invariant to diagonal rescaling of gradients and is well suited for problems that are large in terms of dataparameters args learningrate a float a keras optimizers schedules learningrateschedule instance or a callable that takes no arguments and returns the actual value to use the learning rate defaults to 0 001 beta1 a float value or a constant float tensor or a callable that takes no arguments and returns the actual value to use the exponential decay rate for the 1st moment estimates defaults to 0 9 beta2 a float value or a constant float tensor or a callable that takes no arguments and returns the actual value to use the exponential decay rate for the 2nd moment estimates defaults to 0 999 epsilon a small constant for numerical stability this epsilon is epsilon hat in the kingma and ba paper in the formula just before section 2 1 not the epsilon in algorithm 1 of the paper defaults to 1e7 amsgrad boolean whether to apply amsgrad variant of this algorithm from the paper on the convergence of adam and beyond defaults to false baseoptimizerkeywordargs initialize optimizer variables adam optimizer has 3 types of variables momentums velocities and velocityhat only set when amsgrad is applied args varlist list of model variables to build adam variables on update step given gradient and the associated model variable lr ops castlearningrate variable dtype gradient ops castgradient variable dtype localstep ops castself iterations 1 variable dtype beta1power ops power ops castself beta1 variable dtype localstep beta2power ops power ops castself beta2 variable dtype localstep m self momentumsself getvariableindexvariable v self velocitiesself getvariableindexvariable alpha lr ops sqrt1 beta2power 1 beta1power self assignadd m ops multiplyops subtractgradient m 1 self beta1 self assignadd v ops multiply ops subtractops squaregradient v 1 self beta2 if self amsgrad vhat self velocityhatsself getvariableindexvariable self assignvhat ops maximumvhat v v vhat self assignsub variable ops divide ops multiplym alpha ops addops sqrtv self epsilon def getconfigself config super getconfig config update beta1 self beta1 beta2 self beta2 epsilon self epsilon amsgrad self amsgrad return config adam doc adam doc replace baseoptimizerkeywordargs optimizer baseoptimizerkeywordargs optimizer that implements the adam algorithm adam optimization is a stochastic gradient descent method that is based on adaptive estimation of first order and second order moments according to kingma et al 2014 http arxiv org abs 1412 6980 the method is computationally efficient has little memory requirement invariant to diagonal rescaling of gradients and is well suited for problems that are large in terms of data parameters args learning_rate a float a keras optimizers schedules learningrateschedule instance or a callable that takes no arguments and returns the actual value to use the learning rate defaults to 0 001 beta_1 a float value or a constant float tensor or a callable that takes no arguments and returns the actual value to use the exponential decay rate for the 1st moment estimates defaults to 0 9 beta_2 a float value or a constant float tensor or a callable that takes no arguments and returns the actual value to use the exponential decay rate for the 2nd moment estimates defaults to 0 999 epsilon a small constant for numerical stability this epsilon is epsilon hat in the kingma and ba paper in the formula just before section 2 1 not the epsilon in algorithm 1 of the paper defaults to 1e 7 amsgrad boolean whether to apply amsgrad variant of this algorithm from the paper on the convergence of adam and beyond defaults to false base_optimizer_keyword_args initialize optimizer variables adam optimizer has 3 types of variables momentums velocities and velocity_hat only set when amsgrad is applied args var_list list of model variables to build adam variables on update step given gradient and the associated model variable
from keras import ops
from keras.api_export import keras_export
from keras.optimizers import optimizer


@keras_export(["keras.optimizers.Adam"])
class Adam(optimizer.Optimizer):
    

    def __init__(
        self,
        learning_rate=0.001,
        beta_1=0.9,
        beta_2=0.999,
        epsilon=1e-7,
        amsgrad=False,
        weight_decay=None,
        clipnorm=None,
        clipvalue=None,
        global_clipnorm=None,
        use_ema=False,
        ema_momentum=0.99,
        ema_overwrite_frequency=None,
        name="adam",
        **kwargs,
    ):
        super().__init__(
            learning_rate=learning_rate,
            name=name,
            weight_decay=weight_decay,
            clipnorm=clipnorm,
            clipvalue=clipvalue,
            global_clipnorm=global_clipnorm,
            use_ema=use_ema,
            ema_momentum=ema_momentum,
            ema_overwrite_frequency=ema_overwrite_frequency,
            **kwargs,
        )
        self.beta_1 = beta_1
        self.beta_2 = beta_2
        self.epsilon = epsilon
        self.amsgrad = amsgrad

    def build(self, var_list):
        
        if self.built:
            return
        super().build(var_list)
        self._momentums = []
        self._velocities = []
        for var in var_list:
            self._momentums.append(
                self.add_variable_from_reference(
                    reference_variable=var, name="momentum"
                )
            )
            self._velocities.append(
                self.add_variable_from_reference(
                    reference_variable=var, name="velocity"
                )
            )
        if self.amsgrad:
            self._velocity_hats = []
            for var in var_list:
                self._velocity_hats.append(
                    self.add_variable_from_reference(
                        reference_variable=var, name="velocity_hat"
                    )
                )

    def update_step(self, gradient, variable, learning_rate):
        
        lr = ops.cast(learning_rate, variable.dtype)
        gradient = ops.cast(gradient, variable.dtype)
        local_step = ops.cast(self.iterations + 1, variable.dtype)
        beta_1_power = ops.power(
            ops.cast(self.beta_1, variable.dtype), local_step
        )
        beta_2_power = ops.power(
            ops.cast(self.beta_2, variable.dtype), local_step
        )

        m = self._momentums[self._get_variable_index(variable)]
        v = self._velocities[self._get_variable_index(variable)]

        alpha = lr * ops.sqrt(1 - beta_2_power) / (1 - beta_1_power)

        self.assign_add(
            m, ops.multiply(ops.subtract(gradient, m), 1 - self.beta_1)
        )
        self.assign_add(
            v,
            ops.multiply(
                ops.subtract(ops.square(gradient), v), 1 - self.beta_2
            ),
        )
        if self.amsgrad:
            v_hat = self._velocity_hats[self._get_variable_index(variable)]
            self.assign(v_hat, ops.maximum(v_hat, v))
            v = v_hat
        self.assign_sub(
            variable,
            ops.divide(
                ops.multiply(m, alpha), ops.add(ops.sqrt(v), self.epsilon)
            ),
        )

    def get_config(self):
        config = super().get_config()
        config.update(
            {
                "beta_1": self.beta_1,
                "beta_2": self.beta_2,
                "epsilon": self.epsilon,
                "amsgrad": self.amsgrad,
            }
        )
        return config


Adam.__doc__ = Adam.__doc__.replace(
    "{{base_optimizer_keyword_args}}", optimizer.base_optimizer_keyword_args
)