Model Training and Callbacks

Relevant source files

• tensorflow/python/keras/activations.py
• tensorflow/python/keras/backend.py
• tensorflow/python/keras/backend_config.py
• tensorflow/python/keras/callbacks.py
• tensorflow/python/keras/callbacks_v1.py
• tensorflow/python/keras/constraints.py
• tensorflow/python/keras/distribute/distributed_training_utils_v1.py
• tensorflow/python/keras/engine/base_layer.py
• tensorflow/python/keras/engine/base_layer_v1.py
• tensorflow/python/keras/engine/compile_utils.py
• tensorflow/python/keras/engine/data_adapter.py
• tensorflow/python/keras/engine/functional.py
• tensorflow/python/keras/engine/input_layer.py
• tensorflow/python/keras/engine/keras_tensor.py
• tensorflow/python/keras/engine/training.py
• tensorflow/python/keras/engine/training_arrays_v1.py
• tensorflow/python/keras/engine/training_distributed_v1.py
• tensorflow/python/keras/engine/training_eager_v1.py
• tensorflow/python/keras/engine/training_generator_v1.py
• tensorflow/python/keras/engine/training_utils.py
• tensorflow/python/keras/engine/training_utils_v1.py
• tensorflow/python/keras/engine/training_v1.py
• tensorflow/python/keras/keras_parameterized.py
• tensorflow/python/keras/layers/advanced_activations.py
• tensorflow/python/keras/layers/convolutional.py
• tensorflow/python/keras/layers/convolutional_recurrent.py
• tensorflow/python/keras/layers/merge.py
• tensorflow/python/keras/layers/recurrent.py
• tensorflow/python/keras/losses.py
• tensorflow/python/keras/metrics.py
• tensorflow/python/keras/mixed_precision/loss_scale_optimizer.py
• tensorflow/python/keras/mixed_precision/policy.py
• tensorflow/python/keras/models.py
• tensorflow/python/keras/optimizer_v1.py
• tensorflow/python/keras/saving/hdf5_format.py
• tensorflow/python/keras/saving/model_config.py
• tensorflow/python/keras/saving/save.py
• tensorflow/python/keras/saving/saved_model/base_serialization.py
• tensorflow/python/keras/saving/saved_model/json_utils.py
• tensorflow/python/keras/saving/saved_model/layer_serialization.py
• tensorflow/python/keras/saving/saved_model/load.py
• tensorflow/python/keras/saving/saved_model/metric_serialization.py
• tensorflow/python/keras/saving/saved_model/save.py
• tensorflow/python/keras/saving/saved_model/save_impl.py
• tensorflow/python/keras/saving/saved_model/serialized_attributes.py
• tensorflow/python/keras/saving/saved_model/utils.py
• tensorflow/python/keras/saving/saving_utils.py
• tensorflow/python/keras/testing_utils.py
• tensorflow/python/keras/utils/control_flow_util.py
• tensorflow/python/keras/utils/conv_utils.py
• tensorflow/python/keras/utils/data_utils.py
• tensorflow/python/keras/utils/generic_utils.py
• tensorflow/python/keras/utils/io_utils.py
• tensorflow/python/keras/utils/losses_utils.py
• tensorflow/python/keras/utils/metrics_utils.py
• tensorflow/python/keras/utils/tf_utils.py
• tensorflow/python/training/experimental/loss_scale.py
• tensorflow/python/training/experimental/mixed_precision.py
• tensorflow/python/training/experimental/mixed_precision_test.py
• tensorflow/tools/api/golden/v2/tensorflow.__internal__.mixed_precision.pbtxt

The Keras training infrastructure in TensorFlow provides a high-level API for orchestrating machine learning workflows. It centers around the Model.fit execution loop, which abstracts data ingestion via adapters, loss and metric computation, and a sophisticated callback system for monitoring and side-effect management.

1. The Training Loop (Model.fit)

The Model.fit method in tensorflow/python/keras/engine/training.py is the entry point for supervised training. It handles the transition from raw input data to executed gradient updates.

Data Adaptation

Before training begins, the DataAdapter system converts various input types (NumPy arrays, tf.data.Dataset, Python generators) into a unified tf.data.Dataset format tensorflow/python/keras/engine/data_adapter.py55-80

Adapter Class	Input Type Handled	Source
TensorDataAdapter	Eager Tensors	tensorflow/python/keras/engine/data_adapter.py293
NumpyArrayDataAdapter	NumPy Arrays	tensorflow/python/keras/engine/data_adapter.py195
DatasetAdapter	tf.data.Dataset	tensorflow/python/keras/engine/data_adapter.py535
GeneratorDataAdapter	Python Generators	tensorflow/python/keras/engine/data_adapter.py787

Execution Flow

The training loop utilizes train_function, which is a compiled tf.function that wraps a single step of training tensorflow/python/keras/engine/training.py821-840

Model Training Data Flow

Sources: tensorflow/python/keras/engine/training.py759-800 tensorflow/python/keras/engine/data_adapter.py55-138

2. Callback System

Callbacks are objects that can perform actions at various stages of training (e.g., at the start/end of epochs or batches). The CallbackList class acts as a container and dispatcher for these events tensorflow/python/keras/callbacks.py195-200

Key Lifecycle Methods

• on_train_begin / on_train_end: Called at the start and end of fit.
• on_epoch_begin / on_epoch_end: Called at the start and end of every epoch.
• on_train_batch_begin / on_train_batch_end: Called before and after every batch update.

Built-in Callbacks

• ModelCheckpoint: Periodically saves the Keras model or weights to a given path tensorflow/python/keras/callbacks.py1218-1250
• EarlyStopping: Monitors a specific metric and stops training when improvement ceases tensorflow/python/keras/callbacks.py1655-1680
• TensorBoard: Writes logs for visualization in the TensorBoard dashboard tensorflow/python/keras/callbacks.py2058-2100
• History: Automatically applied to all models to record metrics for each epoch tensorflow/python/keras/callbacks.py1100-1110

Sources: tensorflow/python/keras/callbacks.py17-133

3. Metrics and Loss Computation

Metrics and losses are specialized Layer subclasses that track state across batches.

Metrics API

The Metric base class requires subclasses to implement update_state(), result(), and reset_state() tensorflow/python/keras/metrics.py74-144

• update_state: Logic for accumulating statistics (e.g., sum of squares, true positive counts) tensorflow/python/keras/metrics.py129-140
• result: Logic for computing the final metric value from accumulated state tensorflow/python/keras/metrics.py141-142

Loss API

The Loss base class handles the reduction of per-sample losses into a scalar value tensorflow/python/keras/losses.py48-84 It supports different Reduction types: AUTO, SUM_OVER_BATCH_SIZE, SUM, and NONE tensorflow/python/keras/losses.py86-103

Metric State Management

Sources: tensorflow/python/keras/metrics.py74-144 tensorflow/python/keras/losses.py48-157

4. Mixed Precision Training

Keras supports mixed precision training, where most operations use float16 for speed, while certain variables (like weights) remain in float32 for numerical stability.

Implementation Details

• Policy: Controls the compute_dtype and variable_dtype for layers tensorflow/python/keras/mixed_precision/policy.py30-60
• LossScaleOptimizer: Wraps a standard optimizer to scale the loss. This prevents "gradient underflow" where small gradients vanish when converted to float16 tensorflow/python/keras/mixed_precision/loss_scale_optimizer.py54-100
• Autocasting: The Layer class automatically casts inputs to the policy's compute_dtype during the call() method tensorflow/python/keras/engine/base_layer.py930-950

Mixed Precision Entity Mapping

Code Entity	Responsibility	Source
Policy	Defines dtype pairs (e.g., "mixed_float16")	tensorflow/python/keras/mixed_precision/policy.py30
LossScaleOptimizer	Scales loss before backprop, unscales gradients after	tensorflow/python/keras/mixed_precision/loss_scale_optimizer.py54
autocast_variable	Ensures variables are read in the correct dtype	tensorflow/python/keras/engine/base_layer.py53

Sources: tensorflow/python/keras/mixed_precision/policy.py30-60 tensorflow/python/keras/mixed_precision/loss_scale_optimizer.py54-150 tensorflow/python/keras/engine/base_layer.py97-150