Training, saving and serving

model.fit with callbacks, a custom GradientTape loop, the tf.data pipeline, and exporting a model whose preprocessing travels with it.

Training with fit and callbacks

import tensorflow as tf
from tensorflow import keras

callbacks = [
    keras.callbacks.EarlyStopping(
        monitor="val_loss", patience=3, restore_best_weights=True),
    keras.callbacks.ModelCheckpoint(
        "best.keras", monitor="val_loss", save_best_only=True),
    keras.callbacks.ReduceLROnPlateau(
        monitor="val_loss", factor=0.5, patience=2, min_lr=1e-6),
    keras.callbacks.TensorBoard(log_dir="logs/", histogram_freq=1),
]

history = model.fit(
    x_train, y_train,
    validation_split=0.15,        # or validation_data=(x_val, y_val)
    epochs=50,
    batch_size=64,
    callbacks=callbacks,
    verbose=2,
)

history.history.keys()            # loss, acc, val_loss, val_acc

restore_best_weights=True matters: without it, early stopping leaves the model holding the weights from the last epoch, which may be worse than the best one seen a few epochs earlier.

A custom training loop

optimizer = keras.optimizers.Adam(1e-3)
loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=True)
train_acc = keras.metrics.SparseCategoricalAccuracy()

@tf.function
def train_step(x, y):
    with tf.GradientTape() as tape:
        logits = model(x, training=True)     # training=True enables dropout
        loss = loss_fn(y, logits)
    grads = tape.gradient(loss, model.trainable_variables)
    grads, _ = tf.clip_by_global_norm(grads, 1.0)
    optimizer.apply_gradients(zip(grads, model.trainable_variables))
    train_acc.update_state(y, logits)
    return loss

dataset = (tf.data.Dataset.from_tensor_slices((x_train, y_train))
           .shuffle(10_000)
           .batch(64)
           .prefetch(tf.data.AUTOTUNE))

for epoch in range(5):
    for x_batch, y_batch in dataset:
        loss = train_step(x_batch, y_batch)
    print(epoch, float(loss), float(train_acc.result()))
    train_acc.reset_state()

• Use a custom loop when you need two optimisers, gradient penalty terms, or per-sample weighting that fit cannot express.
• training=True is not optional: it is what makes dropout and batch normalisation behave correctly at training time.
• .prefetch(tf.data.AUTOTUNE) overlaps data loading with the GPU and is the cheapest large speedup available.
• Wrap the step in @tf.function, pass tensors (not Python numbers) as arguments, and the loop traces once instead of per step.

Saving and serving

model.save("model.keras")                # Keras v3 format: architecture, weights, state
model.save("saved_model_dir")            # SavedModel, for TF Serving

restored = keras.models.load_model("model.keras")
restored.predict(x_test[:8])

latest = keras.models.load_model("best.keras")   # written by ModelCheckpoint
latest.evaluate(x_test, y_test)                  # evaluate, not predict, for metrics

keras.backend.clear_session()            # free GPU memory between experiments

# put preprocessing INSIDE the model so serving matches training exactly
normalizer = keras.layers.Normalization(axis=-1)
normalizer.adapt(x_train)                # learns mean and variance from train only

model = keras.Sequential([
    keras.Input(shape=x_train.shape[1:]),
    normalizer,
    keras.layers.Dense(64, activation="relu"),
    keras.layers.Dense(1, activation="sigmoid"),
])
model.compile(optimizer="adam", loss="binary_crossentropy", metrics=["accuracy"])

FAQ

Related

Building a Keras model Pipelines and saving models

Last refreshed 2026-09-18.