Optimizers

from tensorflow.keras import optimizers

# Pass as string (default hyperparameters)
model.compile(optimizer='adam')

# Pass as instance (custom hyperparameters)
opt = optimizers.Adam(learning_rate=1e-3, beta_1=0.9, beta_2=0.999, epsilon=1e-7)
model.compile(optimizer=opt)

# Apply gradients manually
grads = tape.gradient(loss, model.trainable_variables)
opt.apply_gradients(zip(grads, model.trainable_variables))

# Inspect / update LR at runtime
opt.learning_rate.assign(1e-4)
float(opt.learning_rate)
opt.iterations       # number of gradient steps taken

optimizers.SGD(
    learning_rate=0.01,
    momentum=0.0,          # 0.9 typical with momentum
    nesterov=False,        # True = NAG
    weight_decay=0.0,
    clipnorm=None,
    clipvalue=None,
)

# With Nesterov momentum
optimizers.SGD(learning_rate=0.01, momentum=0.9, nesterov=True)

optimizers.Adam(
    learning_rate=0.001,
    beta_1=0.9,            # first-moment decay
    beta_2=0.999,          # second-moment decay
    epsilon=1e-7,          # numerical stability (try 1e-5 for transformers)
    amsgrad=False,         # use AMSGrad variant
    weight_decay=0.0,
    clipnorm=None,
    clipvalue=None,
    global_clipnorm=None,
)

opt = optimizers.AdamW(
    learning_rate=1e-3,
    weight_decay=1e-4,     # decoupled weight decay (not L2 regularization)
    beta_1=0.9,
    beta_2=0.999,
    epsilon=1e-7,
)

# Skip decay for biases / norm params — a method, NOT a constructor kwarg.
# Call before the first training step (before the optimizer builds).
opt.exclude_from_weight_decay(var_names=['bias', 'gamma', 'beta'])
# or pass variables directly:
opt.exclude_from_weight_decay(var_list=[layer.bias for layer in dense_layers])

optimizers.RMSprop(
    learning_rate=0.001,
    rho=0.9,               # discounting factor for running average
    momentum=0.0,
    epsilon=1e-7,
    centered=False,        # True = normalize by estimated variance
)

# Clip by global L2 norm (recommended — preserves gradient direction)
opt = optimizers.Adam(learning_rate=1e-3, global_clipnorm=1.0)

# Clip by per-parameter L2 norm
opt = optimizers.Adam(clipnorm=1.0)

# Clip by value (clips each component independently)
opt = optimizers.Adam(clipvalue=0.5)

# Manual in custom loop
grads, global_norm = tf.clip_by_global_norm(grads, clip_norm=1.0)

# Cosine decay with warm restarts
schedule = optimizers.schedules.CosineDecayRestarts(
    initial_learning_rate=1e-2,
    first_decay_steps=1000,
    t_mul=2.0,
    m_mul=1.0,
    alpha=0.0,
)
opt = optimizers.Adam(learning_rate=schedule)

# Warmup + cosine (manual using LambdaLR equivalent)
class WarmupCosineSchedule(optimizers.schedules.LearningRateSchedule):
    def __init__(self, warmup_steps, total_steps, peak_lr, min_lr=0.0):
        self.warmup_steps = warmup_steps
        self.total_steps  = total_steps
        self.peak_lr      = peak_lr
        self.min_lr       = min_lr

    def __call__(self, step):
        step = tf.cast(step, tf.float32)
        warmup = step / self.warmup_steps * self.peak_lr
        decay  = self.min_lr + 0.5 * (self.peak_lr - self.min_lr) * (
            1 + tf.cos(tf.constant(3.14159) * (step - self.warmup_steps)
                       / (self.total_steps - self.warmup_steps)))
        return tf.where(step < self.warmup_steps, warmup, decay)

    def get_config(self):
        return {'warmup_steps': self.warmup_steps, 'total_steps': self.total_steps,
                'peak_lr': self.peak_lr, 'min_lr': self.min_lr}

import tensorflow as tf
from tensorflow import keras

keras.mixed_precision.set_global_policy('mixed_float16')

# compile() auto-wraps the optimizer under the mixed_float16 policy;
# wrap manually only for custom training loops
base_opt = keras.optimizers.Adam(1e-3)
opt = keras.mixed_precision.LossScaleOptimizer(base_opt)

# In custom loop with mixed precision
with tf.GradientTape() as tape:
    logits = model(x, training=True)
    loss   = loss_fn(y, logits)
    scaled = opt.scale_loss(loss)

grads = tape.gradient(scaled, model.trainable_variables)
opt.apply(grads, model.trainable_variables)  # unscales automatically

opt = optimizers.Adam(1e-3)

# After at least one step:
opt.variables                # all state variables (momentums, velocities, ...)
opt.iterations               # step counter

# Optimizer state is saved/restored with the model:
model.save('model.keras')    # .keras files include optimizer state
model.save_weights('ckpt.weights.h5')  # weights only, no optimizer state

# Zero gradients (rare in TF; done automatically between steps)

# Exponential moving average of weights (evaluate/save with smoothed weights)
opt = optimizers.AdamW(
    learning_rate=1e-3,
    weight_decay=1e-4,
    use_ema=True,
    ema_momentum=0.999,
    ema_overwrite_frequency=None,  # int N: swap EMA weights into the model every N steps
)
model.compile(optimizer=opt, loss='mse')
# With ema_overwrite_frequency=None, finalize_variable_values() copies the
# EMA values into the model at the end of fit()/evaluate() automatically.

# Lion — sign-based, memory-efficient (one moment buffer instead of two)
optimizers.Lion(learning_rate=1e-4, beta_1=0.9, beta_2=0.99, weight_decay=0.0)

# Lamb — layer-wise adaptive rates for very large batches
optimizers.Lamb(learning_rate=1e-3)