iwantcoding.com
🔥 Daily 👥 Rooms 🏆 Top Log in Sign up

Optimizers

By · Reviewed & last updated · Part of TensorFlow Tutorial

Optimisers update weights to minimise loss. Adam is the modern default; SGD + momentum still wins on big vision models; AdamW is the standard for transformers. Pick based on the problem, not by reflex.

Configure, schedule, clip

EXAMPLE
import tensorflow as tf
from tensorflow.keras import optimizers, callbacks

# 1) Default picks
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')
model.compile(optimizer=optimizers.Adam(learning_rate=1e-3))
model.compile(optimizer=optimizers.AdamW(learning_rate=3e-4, weight_decay=1e-2))
model.compile(optimizer=optimizers.SGD(learning_rate=0.1, momentum=0.9, nesterov=True))
model.compile(optimizer=optimizers.RMSprop(learning_rate=1e-3))

# 2) Per-layer learning rates — different rates for backbone vs head
backbone_vars = [v for v in model.trainable_variables if 'backbone' in v.name]
head_vars     = [v for v in model.trainable_variables if 'head'     in v.name]

backbone_opt = optimizers.AdamW(1e-4, weight_decay=1e-2)
head_opt     = optimizers.AdamW(1e-3, weight_decay=1e-2)

@tf.function
def train_step(xb, yb):
    with tf.GradientTape() as tape:
        loss = loss_fn(yb, model(xb, training=True))
    grads = tape.gradient(loss, model.trainable_variables)
    bb = grads[:len(backbone_vars)]
    hd = grads[len(backbone_vars):]
    backbone_opt.apply_gradients(zip(bb, backbone_vars))
    head_opt.apply_gradients(zip(hd, head_vars))
    return loss

# 3) Learning-rate schedule — cosine decay with warmup
schedule = optimizers.schedules.CosineDecay(
    initial_learning_rate=1e-3,
    decay_steps=10_000,
    alpha=0.01,
)
model.compile(optimizer=optimizers.AdamW(learning_rate=schedule, weight_decay=1e-2))

# Or: built-in callback
model.fit(Xtr, ytr,
    callbacks=[
        callbacks.LearningRateScheduler(
            lambda epoch, lr: lr * 0.1 if epoch in {5, 10} else lr,
        ),
        callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.5, patience=3),
    ],
)

# 4) Gradient clipping — stops loss explosions
model.compile(optimizer=optimizers.Adam(learning_rate=1e-3, clipnorm=1.0))
#   clipvalue=N → clip per-coordinate
#   clipnorm=N  → clip global gradient norm (preferred)

# 5) Pick a default by task
#   Vision CNNs / ResNets        — SGD + momentum 0.9, cosine schedule
#   Transformers / NLP / LLMs     — AdamW, β₁=0.9, β₂=0.95, weight_decay=0.01-0.1, warmup
#   Tabular MLPs                  — Adam, lr=1e-3
#   Hard losses / sparse features — Adam, then RMSprop
#   Embeddings / SGDClassifier    — SGD

Why it matters

For most modern models, AdamW + cosine decay + 3-5% linear warmup + gradient clipping is the recipe. Get those four right and you usually outrun “more compute”.

Tip: Tweak the snippet with Try it Yourself », then sit the quiz at the bottom of the page.

Example

Example
from tensorflow.keras import optimizers
opt = optimizers.AdamW(learning_rate=1e-3, weight_decay=1e-4)
Try it Yourself »

Discussion

Loading…