TensorFlow
Keras
Functional API
training flag
deep learning

Passing trainingtrue when using Tensorflow 2's Keras Functional API

Master System Design with Codemia

Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.

Introduction

In TensorFlow 2.x, the training flag tells certain layers whether they should behave in training mode or inference mode. In a normal Keras Functional API model, you usually do not hardcode training=True while building the graph. Instead, you let Keras choose the correct mode through fit, evaluate, and predict.

Which Layers Care About training

Most layers ignore the training flag. A few important ones do not.

  • 'Dropout drops activations during training and passes everything through during inference.'
  • 'BatchNormalization updates moving statistics during training and uses stored statistics during inference.'

That is why the flag matters. If those layers run in the wrong mode, the model can behave incorrectly even though the code still executes.

The Normal Functional API Pattern

In ordinary Functional API code, build the graph without forcing training mode.

python
1import tensorflow as tf
2
3inputs = tf.keras.Input(shape=(16,))
4x = tf.keras.layers.Dense(32, activation="relu")(inputs)
5x = tf.keras.layers.Dropout(0.3)(x)
6outputs = tf.keras.layers.Dense(1)(x)
7
8model = tf.keras.Model(inputs, outputs)
9model.compile(optimizer="adam", loss="mse")

With this model:

  • 'model.fit(...) uses training behavior,'
  • 'model.evaluate(...) uses inference behavior for relevant layers,'
  • 'model.predict(...) uses inference behavior.'

That automatic mode switching is what you want most of the time.

When You Should Forward training

If you write a custom layer or custom model that internally uses dropout, batch normalization, or any other mode-sensitive sublayer, accept training in call and pass it through.

python
1import tensorflow as tf
2
3
4class Encoder(tf.keras.layers.Layer):
5    def __init__(self):
6        super().__init__()
7        self.dense = tf.keras.layers.Dense(32, activation="relu")
8        self.dropout = tf.keras.layers.Dropout(0.3)
9
10    def call(self, inputs, training=None):
11        x = self.dense(inputs)
12        return self.dropout(x, training=training)
13
14
15inputs = tf.keras.Input(shape=(16,))
16x = Encoder()(inputs)
17outputs = tf.keras.layers.Dense(1)(x)
18model = tf.keras.Model(inputs, outputs)

This is the correct place to deal with training: you are relaying Keras's mode decision, not overriding it globally.

Why Hardcoding training=True Is Usually Wrong

You can force training mode directly:

python
1import tensorflow as tf
2
3inputs = tf.keras.Input(shape=(16,))
4x = tf.keras.layers.Dense(32, activation="relu")(inputs)
5x = tf.keras.layers.Dropout(0.3)(x, training=True)
6outputs = tf.keras.layers.Dense(1)(x)
7
8model = tf.keras.Model(inputs, outputs)

But now dropout stays active even during inference. That is usually a bug, because predictions become stochastic and systematically different from the intended deployment behavior.

The same kind of problem applies to batch normalization, where forcing training mode changes whether moving statistics are updated and whether stored statistics are used.

Legitimate Exceptions

There are special cases where forcing training behavior is intentional. A common example is Monte Carlo dropout, where you deliberately keep dropout active during inference to estimate uncertainty.

That is an advanced technique, not the default architecture pattern. If you are not explicitly trying to do that, do not hardcode training=True.

Custom Training Loops Follow the Same Rule

If you use a custom loop instead of model.fit, choose the mode at the call site:

python
1with tf.GradientTape() as tape:
2    y_pred = model(batch_x, training=True)
3    loss = loss_fn(batch_y, y_pred)
4
5val_pred = model(val_x, training=False)

This keeps the training-versus-inference decision explicit and avoids leaking training behavior into validation or serving.

Common Pitfalls

A common mistake is hardcoding training=True while building the Functional API graph and then wondering why predictions are unstable. Dropout and batch normalization are often the reason.

Another issue is forgetting to accept and forward training in custom layers that contain stateful sublayers.

Developers also sometimes assume every layer cares about the flag. Most do not, so the important question is whether the specific sublayers inside your model are mode-sensitive.

Summary

  • In normal Functional API code, do not force training=True during model construction.
  • Let Keras manage training and inference mode through its standard APIs.
  • Forward training explicitly inside custom call methods when your layer contains mode-sensitive sublayers.
  • Hardcoding training mode is only appropriate for special cases such as deliberate Monte Carlo dropout.
  • Understanding which layers depend on training prevents subtle model bugs.

Course illustration
Course illustration

All Rights Reserved.