What training.py Actually Does

training.py defines the core tf.keras.Model engine: compile, fit, evaluate, predict, and the machinery they rely on—data adapters, metrics, distribution, tf.function wrapping, callbacks, and saving/loading weights.

Project: tensorflow/tensorflow

tensorflow/
  python/
    keras/
      engine/
        base_layer.py
        compile_utils.py
        data_adapter.py
        training.py   <-- this file (Model training engine)
      saving/
        save.py
        hdf5_format.py
      utils/
        tf_utils.py
    distribute/
      collective_all_reduce_strategy.py

Call graph (simplified):

User code
  |
  v
Model.fit()
  -> _assert_compile_was_called()
  -> data_adapter.get_data_handler()
  -> make_train_function()
       -> train_step()
          -> self(x)  # forward pass
          -> compiled_loss / compiled_metrics
          -> optimizer.minimize()
  -> callbacks.CallbackList
  -> optional evaluate() for validation

Conceptually, the file does three things:

• Exposes a friendly facade (compile/fit/evaluate/predict).
• Hides orchestration complexity (data adapters, tf.function, strategies, callbacks, summaries).
• Offers narrow extension points (train_step, test_step, predict_step) for custom logic.

The core design decision is to keep your code focused on “one batch of math” and let the framework own the outer loop. The rest of this article unpacks how that template is built and what we can reuse from it.

The Template Method Behind fit()

The central pattern in training.py is the Template Method: a base class defines a fixed outer algorithm and lets subclasses override individual steps.

For Keras, the outer algorithm is “iterate over batches, integrate with distribution, update metrics, call callbacks”. The overridable step is “what happens for one batch?”, implemented as train_step.

def train_step(self, data):
  """The logic for one training step."""
  data = data_adapter.expand_1d(data)
  x, y, sample_weight = data_adapter.unpack_x_y_sample_weight(data)
  with backprop.GradientTape() as tape:
    y_pred = self(x, training=True)
    loss = self.compiled_loss(
        y, y_pred, sample_weight, regularization_losses=self.losses)
  self.optimizer.minimize(loss, self.trainable_variables, tape=tape)
  self.compiled_metrics.update_state(y, y_pred, sample_weight)
  return_metrics = {}
  for metric in self.metrics:
    result = metric.result()
    if isinstance(result, dict):
      return_metrics.update(result)
    else:
      return_metrics[metric.name] = result
  return return_metrics

This method is intentionally small: unpack the batch, forward pass, loss, gradients, metrics. No tf.function, no distribution logic, no callbacks.

The “how to run this at scale” part lives in make_train_function, which constructs the callable that fit() uses.

def make_train_function(self):
  if self.train_function is not None:
    return self.train_function

  def step_function(model, iterator):
    def run_step(data):
      outputs = model.train_step(data)
      with ops.control_dependencies(_minimum_control_deps(outputs)):
        model._train_counter.assign_add(1)
      return outputs

    data = next(iterator)
    outputs = model.distribute_strategy.run(run_step, args=(data,))
    outputs = reduce_per_replica(
        outputs, self.distribute_strategy, reduction='first')
    write_scalar_summaries(outputs, step=model._train_counter)
    return outputs

  if self._steps_per_execution.numpy().item() == 1:
    def train_function(iterator):
      return step_function(self, iterator)
  else:
    def train_function(iterator):
      for _ in math_ops.range(self._steps_per_execution):
        outputs = step_function(self, iterator)
      return outputs

  if not self.run_eagerly:
    train_function = def_function.function(
        train_function, experimental_relax_shapes=True)
    self.train_tf_function = train_function

  self.train_function = train_function

  if self._cluster_coordinator:
    self.train_function = lambda iterator: self._cluster_coordinator.schedule(
        train_function, args=(iterator,))

  return self.train_function

The layers here are the template in action:

• Inner step: run_step calls your train_step and bumps a counter only on success.
• Distribution wrapper: distribute_strategy.run fans the step out to replicas; reduce_per_replica combines results.
• Execution batching: steps_per_execution unrolls several steps inside one call to reduce Python overhead.
• Graph compilation: optionally wraps everything in tf.function unless run_eagerly is set.
• Cluster scheduling: if a cluster coordinator is present, schedules work on remote workers.

The same pattern appears for test_step/make_test_function and predict_step/make_predict_function. The outer template stays fixed; the inner step is where you customize behavior.

That is the core lesson of the file: keep the per-batch unit tiny and overridable, and wrap it in a framework-controlled shell that handles everything else.

How Distribution Stays Out of Your Way

Distribution strategies complicate almost every aspect of execution, but training.py keeps them out of user code. The main abstraction you see is PerReplica values: one tensor per replica, returned from distribute_strategy.run.

The training engine needs to turn those per-replica values into a single batch result. That’s handled by reduce_per_replica:

def reduce_per_replica(values, strategy, reduction='first'):
  """Reduce PerReplica objects."""

  def _reduce(v):
    if reduction == 'concat' and _collective_all_reduce_multi_worker(strategy):
      return _multi_worker_concat(v, strategy)
    if not _is_per_replica_instance(v):
      return v
    elif reduction == 'first':
      return strategy.unwrap(v)[0]
    elif reduction == 'concat':
      if _is_tpu_multi_host(strategy):
        return _tpu_multi_host_concat(v, strategy)
      else:
        return concat(strategy.unwrap(v))
    else:
      raise ValueError('`reduction` must be "first" or "concat".')

  return nest.map_structure(_reduce, values)

Key ideas here:

• reduction='first' is used when any replica’s value is fine (e.g. scalar losses for logging).
• reduction='concat' is used when you need all examples (e.g. predictions).
• TPU and multi-worker strategies use dedicated helpers to produce a stable global order.

The upside of this design is that distribution behavior is centralized instead of bleeding into every training loop. The downside is the strategy-specific branching inside reduce_per_replica, which grows as new strategies arrive.

The suggested refactor in the analysis is to hide those branches behind a small “strategy reduction adapter”, so reduce_per_replica delegates instead of checking strategy types itself. The principle is the same as with train_step: keep the public template simple and push special cases into narrow, swappable components.

Because training.py owns this reduction point, it’s also a natural place to expose observability hooks such as a “replica skew” metric—how much slower the slowest replica is compared to the fastest. That’s another expression of the same idea: one template, many cross-cutting concerns sitting on top of it.

The Memory Cost of predict()

Training and evaluation keep only the current batch and aggregated metrics in memory. Prediction is different: the API returns all outputs at once, so Model.predict accumulates every batch before concatenating.

outputs = None
with self.distribute_strategy.scope():
  ...
  self.predict_function = self.make_predict_function()
  self._predict_counter.assign(0)
  callbacks.on_predict_begin()
  batch_outputs = None
  for _, iterator in data_handler.enumerate_epochs():
    with data_handler.catch_stop_iteration():
      for step in data_handler.steps():
        callbacks.on_predict_batch_begin(step)
        tmp_batch_outputs = self.predict_function(iterator)
        if data_handler.should_sync:
          context.async_wait()
        batch_outputs = tmp_batch_outputs
        if outputs is None:
          outputs = nest.map_structure(lambda batch_output: [batch_output],
                                       batch_outputs)
        else:
          nest.map_structure_up_to(
              batch_outputs,
              lambda output, batch_output: output.append(batch_output),
              outputs, batch_outputs)
        end_step = step + data_handler.step_increment
        callbacks.on_predict_batch_end(end_step, {'outputs': batch_outputs})
  if batch_outputs is None:
    raise ValueError('Expect x to be a non-empty array or dataset.')
  callbacks.on_predict_end()
all_outputs = nest.map_structure_up_to(batch_outputs, concat, outputs)

The behavior is straightforward but has a clear complexity profile:

Memory usage of predict is proportional to S × B × O (steps × batch size × output size) because all batches are stored until the end.

Practically this means:

• Large inference runs can hit OOM even when training on the same data fits.
• Consumers can’t process early predictions until the entire dataset is done.

The proposed improvement is to extract the loop into a helper that accepts a “collector” callback. That separates “iterate and run predict_function” from “what to do with each batch’s outputs”.

# Illustrative shape, not exact code from the repo

def _accumulate_predictions(self, data_handler, callbacks, collect_fn):
  self.predict_function = self.make_predict_function()
  self._predict_counter.assign(0)
  callbacks.on_predict_begin()
  batch_outputs = None
  for _, iterator in data_handler.enumerate_epochs():
    with data_handler.catch_stop_iteration():
      for step in data_handler.steps():
        callbacks.on_predict_batch_begin(step)
        tmp_batch_outputs = self.predict_function(iterator)
        if data_handler.should_sync:
          context.async_wait()
        batch_outputs = tmp_batch_outputs
        end_step = step + data_handler.step_increment
        collect_fn(batch_outputs, step, end_step)
        callbacks.on_predict_batch_end(end_step, {'outputs': batch_outputs})
  if batch_outputs is None:
    raise ValueError('Expect x to be a non-empty array or dataset.')
  callbacks.on_predict_end()
  return batch_outputs

On top of this template you can implement:

• A default in-memory collector that mimics the current behavior.
• A streaming collector that yields or writes batch outputs incrementally.

Because the memory cost is baked into the current template, the analysis also suggests tracking a metric such as keras.predict_memory_bytes to surface how expensive large prediction runs are. Again, the same theme: a shared loop, plus observability, plus a small customization hook.

Weights and the “God Object” Boundary

The same Model class that owns training also owns saving and loading weights. This is where the “God object” smell shows up: the public facade is necessarily large, and IO concerns sit next to training logic.

Even so, the implementations of save_weights and load_weights are instructive in how they manage formats and contracts.

def save_weights(self,
                 filepath,
                 overwrite=True,
                 save_format=None,
                 options=None):
  ...
  self._assert_weights_created()
  filepath = path_to_string(filepath)
  filepath_is_h5 = saving_utils.is_hdf5_filepath(filepath)
  if save_format is None:
    if filepath_is_h5:
      save_format = 'h5'
    else:
      save_format = 'tf'
  else:
    user_format = save_format.lower().strip()
    if user_format in ('tensorflow', 'tf'):
      save_format = 'tf'
    elif user_format in ('hdf5', 'h5', 'keras'):
      save_format = 'h5'
    else:
      raise ValueError(
          'Unknown format "%s". Was expecting one of {"tf", "h5"}.' % (
              save_format,))
  if save_format == 'tf' and filepath_is_h5:
    raise ValueError(
        ('save_weights got save_format="tf"/"tensorflow", but the '
         'filepath ("%s") looks like an HDF5 file. Omit the ".h5"/".keras" '
         'when saving in TensorFlow format.')
        % filepath)
  ...

Notable practices here:

• Format inference with override: infer format from the extension, but allow explicit override with clear error messages when they conflict.
• Precondition enforcement: _assert_weights_created ensures you don’t silently save an unbuilt model.
• Security hardening: earlier in the file, os.environ.setdefault('HDF5_PLUGIN_PATH', os.devnull) avoids loading arbitrary HDF5 plugins.

load_weights mirrors this with _detect_save_format, distinguishing between HDF5 files, checkpoint prefixes, and SavedModel directories.

The proposed refactor is to move the bodies of save_weights and load_weights into dedicated helpers, so Model remains the facade while IO complexity lives in smaller units.

Design Lessons You Can Reuse

Under the hood of model.fit() is a clear template: a tiny per-batch unit that you can override, wrapped in a heavy-duty shell that handles looping, devices, IO, metrics, and observability. That separation is what makes Keras feel simple while remaining scalable.

1. Isolate the Unit of Work

Define a fixed outer algorithm (looping, logging, retries, distribution) and expose a small hook for “one unit of work” (train_step, test_step, predict_step). Keep that hook free of cross-cutting concerns.

In your own systems, explicitly separate:

• Outer orchestration: resources, concurrency, error boundaries.
• Inner work: one job, one batch, one request.

2. Centralize Cross-Cutting Concerns

Distribution, logging, metrics, and callbacks should not leak into user-defined logic. Keras keeps them in make_* functions and helpers like reduce_per_replica. Even when the internals are complex, they’re at least localized.

When you add environments (new cluster types, backends, or protocols), plug them in behind adapters instead of scattering type checks across the codebase.

3. Make Memory Behavior an API Choice

Model.predict shows how a convenient “give me everything” API quietly enforces an O(N) memory cost. For large workloads, you need a streaming alternative built from the same loop.

• Offer both all-at-once and streaming APIs when results can grow large.
• Attach metrics to the shared loop so you can observe memory and throughput profiles.

4. Enforce Invariants with Explicit Errors

training.py uses guards like _assert_compile_was_called, distribution-scope checks, and _disallow_inside_tf_function to fail fast when contracts are broken. These errors encode assumptions that would otherwise become subtle bugs.

Whenever your design relies on “must be called this way” or “must not run under X”, turn that into a runtime check with a precise message.

5. Keep the Facade, Split the Implementation

Keras’ Model is deliberately a large facade because it’s the primary touchpoint for users. Internally, you can relieve the “God object” pressure by moving training, distribution, and serialization details into focused helpers or mixins.

That combination—a small overridable unit of work, a shared execution template, and extracted helpers for cross-cutting concerns—is what lets model.fit() stay simple while hiding the complexity of running on anything from a single CPU to a large cluster. It’s a pattern you can apply in any non-trivial system where you want power without exposing every sharp edge.