Custom callbacks in TensorFlow - Deep Dive

Overview - Custom callbacks

What is it?

Custom callbacks are special tools you create to watch and react to your machine learning model's training process. They let you add your own actions at different points, like after each training step or at the end of an epoch. This helps you control, monitor, or change training without stopping it. Think of them as personal assistants that keep an eye on your model while it learns.

Why it matters

Without custom callbacks, you would have to stop training to check progress or make changes, which is slow and inefficient. Custom callbacks let you automate tasks like saving the best model, adjusting learning rates, or stopping early when the model is good enough. This makes training smarter, faster, and more reliable, saving time and computing resources.

Where it fits

Before learning custom callbacks, you should understand how model training works in TensorFlow and what callbacks are in general. After mastering custom callbacks, you can explore advanced training techniques like dynamic learning rate schedules, distributed training hooks, or integrating callbacks with TensorBoard for visualization.

Mental Model

Core Idea

A custom callback is a user-defined helper that hooks into the training process to perform actions automatically at key moments.

Think of it like...

Imagine coaching a runner during a race: you watch their pace, cheer them on, give water at certain points, or tell them to slow down if needed. Custom callbacks are like that coach, reacting and helping during the training race of your model.

Training Process Flow
┌───────────────┐
│ Start Training│
└──────┬────────┘
       │
       ▼
┌───────────────┐
│ Epoch begins  │
└──────┬────────┘
       │
       ▼
┌───────────────┐
│ Batch begins  │
└──────┬────────┘
       │
       ▼
┌───────────────┐
│ Batch ends    │
│ (callback)    │
└──────┬────────┘
       │
       ▼
┌───────────────┐
│ Epoch ends    │
│ (callback)    │
└──────┬────────┘
       │
       ▼
┌───────────────┐
│ Training ends │
│ (callback)    │
└───────────────┘

Build-Up - 7 Steps

1

Foundation

What are Callbacks in Training

Concept: Callbacks are objects that let you insert code at specific points during model training.

In TensorFlow, callbacks are classes that have methods called automatically at key events like the start or end of an epoch or batch. They help you monitor training, save models, or change settings without stopping training.

Result

You understand that callbacks are hooks into the training loop that let you run your own code automatically.

Knowing callbacks exist helps you realize training is not a black box; you can interact with it dynamically.

2

Foundation

Using Built-in Callbacks

Concept: TensorFlow provides ready-made callbacks for common tasks like saving models or stopping early.

Examples include ModelCheckpoint to save the best model, EarlyStopping to stop training when no improvement happens, and TensorBoard for visualization. You add them to model.fit() to use.

Result

You can improve training control without writing any new code, just by using existing callbacks.

Seeing built-in callbacks work shows how powerful and flexible the callback system is.

3

Intermediate

Creating a Simple Custom Callback

🤔Before reading on: do you think you can write a callback that prints a message after each epoch? Commit to yes or no.

Concept: You can make your own callback by subclassing tensorflow.keras.callbacks.Callback and overriding methods.

Define a class inheriting from Callback. Override methods like on_epoch_end(self, epoch, logs=None) to run code after each epoch. For example, print the epoch number and loss.

Result

Your callback prints info after every epoch during training.

Understanding subclassing and method overriding unlocks the ability to customize training behavior fully.

4

Intermediate

Accessing Training Data in Callbacks

🤔Before reading on: do you think callbacks can see the model's current loss and accuracy during training? Commit to yes or no.

Concept: Callbacks receive training metrics via the logs dictionary parameter in their methods.

In methods like on_epoch_end, the logs argument contains metrics like loss and accuracy for that epoch. You can use these to make decisions or print summaries.

Result

You can monitor training progress inside your callback and react accordingly.

Knowing how to access metrics lets you build smarter callbacks that adapt to training results.

5

Intermediate

Modifying Training with Callbacks

🤔Before reading on: can callbacks change the learning rate during training? Commit to yes or no.

Concept: Callbacks can modify model parameters or training settings dynamically, like adjusting learning rates.

Inside callbacks, you can access the optimizer and change its learning rate. For example, in on_epoch_end, reduce learning rate if loss plateaus.

Result

Training adapts automatically, potentially improving model performance.

Realizing callbacks can change training parameters opens doors to advanced training strategies.

6

Advanced

Saving Models Conditionally in Callbacks

🤔Before reading on: do you think you can save a model only when validation accuracy improves? Commit to yes or no.

Concept: Custom callbacks can save models based on any condition you define.

In on_epoch_end, check logs for validation accuracy. If it improves over previous best, call model.save() to save the model.

Result

You keep only the best model versions, saving storage and improving deployment readiness.

Knowing how to implement conditional saving prevents wasting resources and ensures you keep the best model.

7

Expert

Handling State and Multiple Events in Callbacks

🤔Before reading on: can a callback keep track of best metrics across epochs and use that in later events? Commit to yes or no.

Concept: Callbacks can maintain internal state and respond to multiple training events for complex logic.

Store variables like best_loss as class attributes. Update them in on_epoch_end. Use on_train_batch_end to log batch-level info. Combine multiple methods for fine control.

Result

Your callback can implement sophisticated monitoring, early stopping, or custom logging.

Understanding stateful callbacks and multiple event hooks enables building robust, production-ready training monitors.

Under the Hood

During training, TensorFlow runs a loop over epochs and batches. At specific points, it calls methods on all callbacks registered. These methods receive context like epoch number and metric logs. The callback methods run synchronously, allowing them to inspect or modify training state. The model and optimizer objects are accessible via the callback's self.model attribute, letting callbacks interact deeply with training internals.

Why designed this way?

Callbacks were designed to separate training logic from monitoring and control code, making training modular and extensible. This avoids cluttering core training code with many special cases. The event-driven design matches natural training milestones, making it intuitive to hook in. Alternatives like polling or manual checks would be inefficient and error-prone.

Training Loop
┌─────────────────────────────┐
│ For each epoch:             │
│  ┌───────────────────────┐ │
│  │ Call on_epoch_begin() │ │
│  └─────────┬─────────────┘ │
│            │               │
│  For each batch:            │
│   ┌──────────────────────┐ │
│   │ Call on_train_batch_  │ │
│   │ begin()               │ │
│   └─────────┬────────────┘ │
│             │              │
│   Train batch              │
│             │              │
│   ┌──────────────────────┐ │
│   │ Call on_train_batch_  │ │
│   │ end()                 │ │
│   └─────────┬────────────┘ │
│             │              │
│  ┌───────────────────────┐ │
│  │ Call on_epoch_end()    │ │
│  └─────────┬─────────────┘ │
│            │               │
│  ┌───────────────────────┐ │
│  │ Call on_train_end()    │ │
│  └───────────────────────┘ │
└─────────────────────────────┘

Myth Busters - 4 Common Misconceptions

Quick: Do callbacks run in parallel with training or synchronously during training steps? Commit to one.

Common Belief: Callbacks run in parallel threads or processes so they don't slow down training.

Reality: Callbacks run synchronously during training steps, meaning they can slow down training if they do heavy work.

Why it matters: If you put slow code in callbacks, training becomes slower, which surprises many who expect callbacks to be lightweight monitors.

Quick: Can callbacks change the model architecture during training? Commit yes or no.

Common Belief: Callbacks can modify the model's layers or structure while training is running.

Reality: Callbacks cannot change the model architecture during training; the model is fixed once training starts.

Why it matters: Trying to change architecture in callbacks leads to errors or undefined behavior, so architecture changes must happen before training.

Quick: Do callbacks automatically save the model after every batch? Commit yes or no.

Common Belief: Callbacks save the model automatically after every batch by default.

Reality: Callbacks only save the model if explicitly programmed to do so; built-in callbacks like ModelCheckpoint save after epochs or conditions.

Why it matters: Assuming automatic saving wastes storage and can cause confusion when models are not saved as expected.

Quick: Can you rely on logs in callbacks to always contain all metrics? Commit yes or no.

Common Belief: The logs dictionary in callbacks always contains every metric you want.

Reality: Logs only contain metrics computed for that step; some metrics may be missing depending on model configuration.

Why it matters: Relying on missing metrics causes bugs or crashes in callbacks that expect them.

Expert Zone

1

Callbacks can access and modify the optimizer's internal state, enabling advanced techniques like custom learning rate schedules or gradient clipping.

2

Callbacks run on the main training thread, so heavy computations or I/O in callbacks can block training and reduce throughput.

3

Stacking multiple callbacks requires careful ordering because some callbacks may depend on others' side effects or modify shared state.

When NOT to use

Custom callbacks are not suitable for changing model architecture or data preprocessing; use model subclassing or data pipelines instead. For very complex training control, consider writing custom training loops with tf.GradientTape for full flexibility.

Production Patterns

In production, callbacks are used to implement early stopping to save resources, dynamic learning rate adjustments for better convergence, and conditional model checkpointing to keep only the best models. They also integrate with monitoring tools like TensorBoard and custom logging systems.

Connections

Observer Pattern (Software Design)

Custom callbacks implement the observer pattern by watching training events and reacting.

Understanding callbacks as observers clarifies how they decouple monitoring logic from core training, a common software design principle.

Event-driven Programming

Callbacks are event handlers triggered by training events like epoch start or batch end.

Recognizing training as an event-driven process helps in designing responsive and modular training workflows.

Sports Coaching

Like a coach giving feedback and adjusting strategy during a game, callbacks monitor and influence training dynamically.

This connection highlights the importance of timely feedback and adaptation in both human and machine learning.

Common Pitfalls

#1 Writing heavy computations inside callback methods causing slow training.

Wrong approach:

class SlowCallback(tf.keras.callbacks.Callback):
def on_epoch_end(self, epoch, logs=None):
import time
time.sleep(5) # Simulate heavy work
print(f"Epoch {epoch} done")

Correct approach:

class FastCallback(tf.keras.callbacks.Callback):
def on_epoch_end(self, epoch, logs=None):
print(f"Epoch {epoch} done")

Root cause:

Misunderstanding that callbacks run synchronously and blocking code delays training.

#2 Trying to change model layers inside a callback during training.

Wrong approach:

class ChangeModelCallback(tf.keras.callbacks.Callback):
def on_epoch_end(self, epoch, logs=None):
self.model.add(tf.keras.layers.Dense(10)) # Invalid during training

Correct approach:

# Modify model architecture before training starts
model = tf.keras.Sequential([...])
model.add(tf.keras.layers.Dense(10))
model.compile(...)
model.fit(...)

Root cause:

Confusing training control with model definition; model structure must be fixed before training.

#3 Assuming logs always contain validation accuracy in on_epoch_end.

Wrong approach:

class LogCallback(tf.keras.callbacks.Callback):
def on_epoch_end(self, epoch, logs=None):
print(logs['val_accuracy']) # KeyError if no validation

Correct approach:

class SafeLogCallback(tf.keras.callbacks.Callback):
def on_epoch_end(self, epoch, logs=None):
val_acc = logs.get('val_accuracy')
if val_acc is not None:
print(val_acc)
else:
print('No validation accuracy available')

Root cause:

Not checking for metric presence in logs causes runtime errors.

Key Takeaways

Custom callbacks let you add your own code to run automatically during model training at key points like batch or epoch ends.

They provide a powerful way to monitor, control, and adapt training without stopping or rewriting the training loop.

Callbacks run synchronously and have access to training metrics and model internals, enabling advanced behaviors like dynamic learning rate changes or conditional saving.

Understanding the event-driven nature and lifecycle of callbacks is essential to use them effectively and avoid common mistakes.

Expert use of callbacks involves managing internal state, combining multiple events, and balancing callback workload to keep training efficient.

Practice

1. What is the main purpose of creating a custom callback in TensorFlow?

easy

A. To run your own code at specific points during model training

B. To change the model architecture automatically

C. To increase the dataset size during training

D. To convert the model to a different format

Solution

1. Step 1: Understand what callbacks do

   Callbacks let you add custom actions during training, like logging or stopping early.

2. Step 2: Identify the role of custom callbacks

   Custom callbacks let you run your own code at points like start/end of epochs or batches.

3. Final Answer:

   To run your own code at specific points during model training -> Option A

4. Quick Check:

   Custom callbacks = run code during training [OK]

Hint: Callbacks run code during training steps [OK]

Common Mistakes:

• Thinking callbacks change model structure
• Confusing callbacks with data augmentation
• Assuming callbacks convert model formats

2. Which method should you override in a custom callback to execute code at the end of each training epoch?

easy

A. on_train_begin

B. on_epoch_end

C. on_batch_begin

D. on_train_end

Solution

1. Step 1: Recall callback method names

   Common callback methods include on_train_begin, on_epoch_end, on_batch_begin, on_train_end.

2. Step 2: Match method to event

   on_epoch_end runs code after each epoch finishes, which is what the question asks.

3. Final Answer:

   on_epoch_end -> Option B

4. Quick Check:

   End of epoch method = on_epoch_end [OK]

Hint: End of epoch callback method is on_epoch_end [OK]

Common Mistakes:

• Using on_train_begin for epoch end
• Confusing batch and epoch methods
• Overriding on_train_end instead

3. What will be printed after training if this custom callback is used?

import tensorflow as tf
class MyCallback(tf.keras.callbacks.Callback):
    def on_epoch_end(self, epoch, logs=None):
        print(f"Epoch {epoch+1} ended with loss {logs['loss']:.2f}")

model = tf.keras.Sequential([tf.keras.layers.Dense(1)])
model.compile(optimizer='sgd', loss='mse')
model.fit([[1.0]], [[2.0]], epochs=2, callbacks=[MyCallback()])

medium

A. Error because logs is None

B. Epoch 0 ended with loss X\nEpoch 1 ended with loss Y

C. No output printed during training

D. Epoch 1 ended with loss X\nEpoch 2 ended with loss Y (X and Y are numbers)

Solution

1. Step 1: Understand on_epoch_end behavior

   The method prints epoch number +1 and loss from logs after each epoch.

2. Step 2: Check training epochs and output

   Training runs 2 epochs, so prints for epoch 1 and 2 with loss values (numbers).

3. Final Answer:

   Epoch 1 ended with loss X\nEpoch 2 ended with loss Y (X and Y are numbers) -> Option D

4. Quick Check:

   Prints after each epoch with loss [OK]

Hint: Epoch count starts at 0, print adds 1 [OK]

Common Mistakes:

• Assuming epoch numbering starts at 1 internally
• Expecting no print output
• Thinking logs is None causing error

4. Identify the error in this custom callback code that tries to stop training when loss is below 0.1:

class StopCallback(tf.keras.callbacks.Callback):
    def on_epoch_end(self, epoch, logs=None):
        if logs['loss'] < 0.1:
            self.model.stop_training = True

model.fit(x_train, y_train, epochs=10, callbacks=[StopCallback()])

medium

A. on_epoch_end method must return a value

B. stop_training should be called as a method

C. logs dictionary might be None causing KeyError

D. Callback class must inherit from tf.keras.Model

Solution

1. Step 1: Check logs usage

   logs can be None if no metrics are passed, so logs['loss'] may cause KeyError.

2. Step 2: Confirm correct way to stop training

   Setting self.model.stop_training = True is correct; no method call needed.

3. Final Answer:

   logs dictionary might be None causing KeyError -> Option C

4. Quick Check:

   Check logs is not None before accessing [OK]

Hint: Always check logs is not None before using keys [OK]

Common Mistakes:

• Calling stop_training as a method
• Expecting on_epoch_end to return a value
• Inheriting from wrong base class

5. You want to save the model weights only when validation accuracy improves during training. Which custom callback method and logic should you implement?

hard

A. Override on_epoch_end, check logs['val_accuracy'], save weights if improved

B. Override on_train_begin, save weights before training starts

C. Override on_batch_end, save weights after every batch

D. Override on_train_end, save weights after training finishes

Solution

1. Step 1: Identify when to check validation accuracy

   Validation accuracy is updated at epoch end, so on_epoch_end is the right method.

2. Step 2: Implement logic to save weights only if val_accuracy improves

   Inside on_epoch_end, compare current val_accuracy with previous best and save if better.

3. Final Answer:

   Override on_epoch_end, check logs['val_accuracy'], save weights if improved -> Option A

4. Quick Check:

   Save on epoch end if val_accuracy improves [OK]

Hint: Use on_epoch_end to monitor validation metrics [OK]

Common Mistakes:

• Saving weights too often (every batch)
• Saving only before or after all training
• Not checking if accuracy improved before saving