TensorFlow: Training Optimization

Optimizers

• Adam: adaptive learning rate per parameter — best default choice for most problems

• SGD + momentum: often better final accuracy than Adam with proper lr schedule

• RMSprop: good for RNNs

• AdamW: Adam with weight decay — better regularization than L2 on Adam

• Learning rate: most impactful hyperparameter. Start with 1e-3, tune with lr finder.

# Learning rate schedules
lr_schedule = keras.optimizers.schedules.CosineDecay(
    initial_learning_rate=1e-3,
    decay_steps=total_steps,
    alpha=1e-5  # minimum lr
)

lr_schedule = keras.optimizers.schedules.ExponentialDecay(
    initial_learning_rate=1e-3,
    decay_steps=1000,
    decay_rate=0.96,
)

model.compile(optimizer=keras.optimizers.Adam(lr_schedule), ...)

Regularization

# L2 regularization
keras.layers.Dense(128, activation='relu',
    kernel_regularizer=keras.regularizers.L2(1e-4))

# Dropout (most common)
keras.layers.Dropout(0.5)   # drop 50% of activations during training

# Batch normalization (before activation)
keras.layers.BatchNormalization()
keras.layers.Dense(128)
keras.layers.Activation('relu')

# Data augmentation (image)
augmentation = keras.Sequential([
    keras.layers.RandomFlip('horizontal'),
    keras.layers.RandomRotation(0.1),
    keras.layers.RandomZoom(0.1),
    keras.layers.RandomContrast(0.1),
])

# Apply during training only
x = augmentation(inputs, training=True)

Custom Training Loop

# For full control over training (research, custom losses)
optimizer = keras.optimizers.Adam(1e-3)
loss_fn = keras.losses.SparseCategoricalCrossentropy()
train_acc_metric = keras.metrics.SparseCategoricalAccuracy()

@tf.function   # compile to graph for speed
def train_step(x_batch, y_batch):
    with tf.GradientTape() as tape:
        logits = model(x_batch, training=True)
        loss = loss_fn(y_batch, logits)

    gradients = tape.gradient(loss, model.trainable_weights)
    optimizer.apply_gradients(zip(gradients, model.trainable_weights))
    train_acc_metric.update_state(y_batch, logits)
    return loss

for epoch in range(epochs):
    for step, (x_batch, y_batch) in enumerate(train_dataset):
        loss = train_step(x_batch, y_batch)

    acc = train_acc_metric.result()
    print(f"Epoch {epoch}: loss={loss:.4f}, acc={acc:.4f}")
    train_acc_metric.reset_state()

GPU & Mixed Precision

# Check GPUs
print(tf.config.list_physical_devices('GPU'))

# Mixed precision (FP16 compute, FP32 weights) — 2-3x faster on modern GPUs
keras.mixed_precision.set_global_policy('mixed_float16')
# Use float32 for output layer:
outputs = keras.layers.Dense(10, dtype='float32')(x)

# Multi-GPU strategy
strategy = tf.distribute.MirroredStrategy()  # single machine, multiple GPUs
with strategy.scope():
    model = build_model()
    model.compile(...)