Train/Validation Split and Early Stopping by Hand in TensorFlow

The Story — Practice Papers and the Real Exam

Splitting a Dataset — and a Hidden Trap

The natural idea is: shuffle, then take some for training and skip them for validation. But look what happens over two epochs.

import tensorflow as tf

ds = tf.data.Dataset.range(10).shuffle(10, seed=42)          # default: reshuffle every epoch!
train, val = ds.take(8), ds.skip(8)

for epoch in range(2):
    t = sorted(int(v) for v in train)
    v = sorted(int(v) for v in val)
    print(f"epoch {epoch}: train {t}  val {v}")

seen_in_train = set()
for _ in range(2): seen_in_train |= {int(v) for v in train}
print("val items also used for training:", sorted(seen_in_train & {int(v) for v in val}))

The Fix: Split Once, Then Shuffle Only the Training Part

full = tf.data.Dataset.range(10).shuffle(10, seed=42, reshuffle_each_iteration=False)   # one fixed order
train = full.take(8).shuffle(8)          # training may reshuffle freely
val = full.skip(8)                        # validation never changes

for epoch in range(2):
    print(f"epoch {epoch}: train {sorted(int(v) for v in train)}  val {sorted(int(v) for v in val)}")

Even simpler: split your arrays (or file lists) before building the datasets, and never mix them again.

Watching Overfitting Happen

We give a big network only 20 noisy points of a sine wave, and keep 100 other points for validation. Then we train for a long time.

import numpy as np
rng = np.random.default_rng(1)

def make(n):
    x = rng.uniform(-3, 3, (n, 1)).astype("float32")
    return x, (np.sin(x) + rng.normal(0, 0.3, x.shape)).astype("float32")

x_train, y_train = make(20)
x_val, y_val = make(100)

def new_model():
    tf.random.set_seed(0)
    return tf.keras.Sequential([
        tf.keras.Input(shape=(1,)),
        tf.keras.layers.Dense(256, activation="relu"),
        tf.keras.layers.Dense(256, activation="relu"),
        tf.keras.layers.Dense(1)])

model = new_model()
opt = tf.keras.optimizers.Adam(0.003)

@tf.function
def train_step():
    with tf.GradientTape() as tape:
        loss = tf.reduce_mean((model(x_train) - y_train) ** 2)
    grads = tape.gradient(loss, model.trainable_variables)
    opt.apply_gradients(zip(grads, model.trainable_variables))
    return loss

def val_loss():
    return float(tf.reduce_mean((model(x_val) - y_val) ** 2))

for epoch in range(1, 401):
    loss = float(train_step())
    if epoch in (10, 40, 80, 200, 400):
        print(f"epoch {epoch:3d}  train loss {loss:.4f}  val loss {val_loss():.4f}")

Train loss keeps falling. Validation loss falls at first, then climbs: after about epoch 40, the extra training only memorises the noise.

Early Stopping by Hand

class EarlyStopping:
    def __init__(self, variables, patience=20, min_delta=1e-4):
        self.variables, self.patience, self.min_delta = variables, patience, min_delta
        self.best, self.best_epoch, self.wait = float("inf"), 0, 0
        self.best_weights = None

    def update(self, epoch, value):
        """Returns True when training should stop."""
        if value < self.best - self.min_delta:
            self.best, self.best_epoch, self.wait = value, epoch, 0
            self.best_weights = [v.numpy().copy() for v in self.variables]   # snapshot
        else:
            self.wait += 1
        return self.wait >= self.patience

    def restore(self):
        for v, w in zip(self.variables, self.best_weights):
            v.assign(w)

model = new_model()
opt = tf.keras.optimizers.Adam(0.003)

@tf.function
def train_step():
    with tf.GradientTape() as tape:
        loss = tf.reduce_mean((model(x_train) - y_train) ** 2)
    grads = tape.gradient(loss, model.trainable_variables)
    opt.apply_gradients(zip(grads, model.trainable_variables))
    return loss

stopper = EarlyStopping(model.trainable_variables, patience=20)
for epoch in range(1, 401):
    train_step()
    v = val_loss()
    if stopper.update(epoch, v):
        print(f"stopped at epoch {epoch}; best was epoch {stopper.best_epoch} (val loss {stopper.best:.4f})")
        break

print(f"val loss before restore: {val_loss():.4f}")
stopper.restore()
print(f"val loss after restore : {val_loss():.4f}")

Try It — Early Stopping Simulator

Golden Rules