TensorFlow Cheatsheet

Callbacks

Use this TensorFlow reference while you build software engineering projects, review code for technical interview prep, or polish examples for a software engineer resume.

Using Callbacks

Pass a list of callback instances to model.fit(). Callbacks fire at epoch, batch, or training boundaries.

from tensorflow.keras import callbacks

model.fit(
    x_train, y_train,
    epochs=100,
    callbacks=[
        callbacks.ModelCheckpoint('best.keras', save_best_only=True),
        callbacks.EarlyStopping(patience=10, restore_best_weights=True),
        callbacks.ReduceLROnPlateau(factor=0.5, patience=5),
        callbacks.TensorBoard(log_dir='logs/'),
    ]
)

Built-in Callbacks — Saving

ModelCheckpoint

callbacks.ModelCheckpoint(
    filepath='checkpoints/model_{epoch:02d}_{val_loss:.4f}.keras',
    monitor='val_loss',          # metric to watch
    verbose=0,
    save_best_only=True,         # only save when metric improves
    save_weights_only=False,     # True saves only weights (not architecture)
    mode='auto',                 # 'min'|'max'|'auto'
    save_freq='epoch',           # 'epoch' or int (save every N batches)
    initial_value_threshold=None,
)

Use .keras extension (Keras 3 native format) or .h5 / SavedModel directory.

EarlyStopping

callbacks.EarlyStopping(
    monitor='val_loss',
    min_delta=1e-4,              # minimum change to count as improvement
    patience=10,                 # epochs with no improvement before stopping
    verbose=1,
    mode='auto',
    baseline=None,               # stop if monitor doesn't beat this value
    restore_best_weights=True,   # roll back to best epoch weights
    start_from_epoch=5,          # ignore early epochs
)

Built-in Callbacks — Training Control

ReduceLROnPlateau

callbacks.ReduceLROnPlateau(
    monitor='val_loss',
    factor=0.5,                  # new_lr = lr * factor
    patience=5,
    verbose=1,
    mode='auto',
    min_delta=1e-4,
    cooldown=0,                  # epochs to wait after reduction before monitoring resumes
    min_lr=1e-7,
)

TensorBoard

callbacks.TensorBoard(
    log_dir='logs/fit/',
    histogram_freq=1,            # compute weight histograms every N epochs (0=off)
    write_graph=True,
    write_images=False,
    write_steps_per_second=True,
    update_freq='epoch',         # 'batch'|'epoch'|int (every N batches)
    profile_batch=0,             # 0=disabled; (2,5) profiles batches 2-5
    embeddings_freq=0,
)

# Launch TensorBoard
# tensorboard --logdir logs/fit

LearningRateScheduler

def lr_schedule(epoch, lr):
    if epoch < 10:
        return lr
    return lr * 0.9

callbacks.LearningRateScheduler(lr_schedule, verbose=1)

Built-in Callbacks — Logging and Utilities

CSVLogger

callbacks.CSVLogger(
    filename='training_log.csv',
    separator=',',
    append=False,   # True to continue logging across sessions
)

BackupAndRestore

callbacks.BackupAndRestore(
    backup_dir='/tmp/backup',
    save_freq='epoch',
    delete_checkpoint=True,   # remove checkpoint after training completes
)

TerminateOnNaN

callbacks.TerminateOnNaN()   # stops training if loss becomes NaN

LambdaCallback

callbacks.LambdaCallback(
    on_epoch_begin=None,
    on_epoch_end=lambda epoch, logs: print(f'Epoch {epoch}: {logs}'),
    on_train_begin=None,
    on_train_end=None,
    on_batch_begin=None,
    on_batch_end=None,
)

Writing Custom Callbacks

class GradientLogger(keras.callbacks.Callback):
    def on_epoch_end(self, epoch, logs=None):
        for layer in self.model.layers:
            for weight in layer.trainable_weights:
                print(f'{weight.name}: norm={tf.norm(weight):.4f}')

class LRWarmup(keras.callbacks.Callback):
    def __init__(self, warmup_epochs, peak_lr):
        super().__init__()
        self.warmup_epochs = warmup_epochs
        self.peak_lr = peak_lr

    def on_epoch_begin(self, epoch, logs=None):
        if epoch < self.warmup_epochs:
            lr = self.peak_lr * (epoch + 1) / self.warmup_epochs
            self.model.optimizer.learning_rate.assign(lr)

class MetricLogger(keras.callbacks.Callback):
    def on_train_begin(self, logs=None):
        self.history = {}

    def on_epoch_end(self, epoch, logs=None):
        for k, v in (logs or {}).items():
            self.history.setdefault(k, []).append(v)

Callback Hooks (all available methods)

Method	Called when
`on_train_begin(logs)`	start of `fit()`
`on_train_end(logs)`	end of `fit()`
`on_epoch_begin(epoch, logs)`	start of each epoch
`on_epoch_end(epoch, logs)`	end of each epoch
`on_train_batch_begin(batch, logs)`	start of each train batch
`on_train_batch_end(batch, logs)`	end of each train batch
`on_test_begin(logs)`	start of `evaluate()`
`on_test_end(logs)`	end of `evaluate()`
`on_test_batch_begin(batch, logs)`	start of each eval batch
`on_test_batch_end(batch, logs)`	end of each eval batch
`on_predict_begin(logs)`	start of `predict()`
`on_predict_end(logs)`	end of `predict()`
`on_predict_batch_begin(batch, logs)`	start of each predict batch
`on_predict_batch_end(batch, logs)`	end of each predict batch

Accessing Model and Optimizer Inside Callback

class MyCallback(keras.callbacks.Callback):
    def on_epoch_end(self, epoch, logs=None):
        model  = self.model
        opt    = model.optimizer
        lr     = float(opt.learning_rate)
        params = self.params   # dict: {'epochs', 'steps', 'verbose', 'metrics'}
        print(f'LR: {lr}')

Using Callbacks in Custom Training Loops

Callbacks can be used manually in custom training loops.

cb_list = keras.callbacks.CallbackList(
    [callbacks.TensorBoard('logs/'), callbacks.EarlyStopping(patience=3)],
    add_history=True,
    add_progbar=True,
    model=model,
)

cb_list.on_train_begin()
for epoch in range(epochs):
    cb_list.on_epoch_begin(epoch)
    logs = {}
    for step, (x, y) in enumerate(train_ds):
        cb_list.on_train_batch_begin(step)
        # ... training step ...
        logs['loss'] = float(loss)
        cb_list.on_train_batch_end(step, logs)
    cb_list.on_epoch_end(epoch, logs)
    if model.stop_training:
        break
cb_list.on_train_end(logs)