Training Loop From Scratch: Epochs

Section 01

The Story — A Full School Year

📖 Real World Analogy

Lessons, Chapters and Exams

A student has a 300-page book. She does not read it all in one sitting. She reads one chapter at a time, and after each chapter she checks her mistakes and adjusts how she thinks. When she finishes the whole book, that is one pass. Then she starts again, a little wiser each time. Now and then she takes a practice exam with questions she has never seen, to check she is really learning and not just memorising.

Training works the same way. A chapter is a batch. One pass through all the data is an epoch. Each batch gives one training step. The practice exam is the validation set, and the score is accuracy.

This is the final lesson of the module. You will put every piece together — tf.Module, your Dense layer, ReLU, softmax cross-entropy and your own Adam — into a complete training loop that learns to classify a tricky spiral dataset.

Data

Split into train and validation. Build a tf.data pipeline: shuffle → batch.

Model

An MLP made of Dense modules with ReLU, ending in logits.

Train step

Forward → loss → gradients → optimizer update. Wrapped in @tf.function.

Epoch loop

Run every batch once. Track average loss and accuracy.

Validate

After each epoch, measure accuracy on data the model never trains on.

Section 02

Step 1 — The Data

Three interleaved spiral arms, one per class. No straight line can separate them, so the network must learn curved borders.

import numpy as np
import tensorflow as tf

def make_spirals(n_per_class=200, classes=3, noise=0.2, seed=0):
    rng = np.random.default_rng(seed)
    X, y = [], []
    for c in range(classes):
        r = np.linspace(0.05, 1.0, n_per_class)
        t = np.linspace(c * 4.0, (c + 1) * 4.0, n_per_class) + rng.normal(0, noise, n_per_class)
        X.append(np.stack([r * np.sin(t), r * np.cos(t)], axis=1))
        y.append(np.full(n_per_class, c))
    X, y = np.concatenate(X).astype("float32"), np.concatenate(y).astype("int32")
    idx = rng.permutation(len(X))
    return X[idx], y[idx]

X, y = make_spirals()
X_train, y_train = X[:480], y[:480]           # 80% to learn from
X_val,   y_val   = X[480:], y[480:]           # 20% practice exam

BATCH = 32
train_ds = (tf.data.Dataset.from_tensor_slices((X_train, y_train))
            .shuffle(len(X_train), seed=1)
            .batch(BATCH))
val_ds = tf.data.Dataset.from_tensor_slices((X_val, y_val)).batch(256)

print("train:", X_train.shape, " val:", X_val.shape)
print("batches per epoch:", len(train_ds), "(last batch has", len(X_train) % BATCH or BATCH, "samples)")

OUTPUT

train: (480, 2) val: (120, 2) batches per epoch: 15 (last batch has 32 samples)

Section 03

Try It — Batches and Epochs

📚 How Many Steps Is One Epoch? Interactive

Set the dataset size, batch size and number of epochs. Each coloured block is one batch (one training step). Press Animate to watch an epoch go by.

samples batch size epochs drop_remainder

Section 04

Step 2 — The Model and the Optimizer

These are the same building blocks you wrote in the earlier lessons, gathered in one place.

import math

class Dense(tf.Module):
    def __init__(self, n_in, n_out, activation=None, name=None):
        super().__init__(name=name)
        std = math.sqrt(2.0 / n_in)                                   # He init
        self.w = tf.Variable(tf.random.normal([n_in, n_out], stddev=std, seed=7), name="w")
        self.b = tf.Variable(tf.zeros([n_out]), name="b")
        self.activation = activation

    def __call__(self, x):
        z = x @ self.w + self.b
        return self.activation(z) if self.activation else z

class MLP(tf.Module):
    def __init__(self, sizes, name=None):
        super().__init__(name=name)
        self.layers = [Dense(a, b, tf.nn.relu if i < len(sizes) - 2 else None)
                       for i, (a, b) in enumerate(zip(sizes[:-1], sizes[1:]))]

    def __call__(self, x):
        for layer in self.layers:
            x = layer(x)
        return x                                                     # logits

class Adam(tf.Module):
    def __init__(self, lr=0.01, b1=0.9, b2=0.999, eps=1e-7):
        super().__init__()
        self.lr, self.b1, self.b2, self.eps = lr, b1, b2, eps
        self.t = tf.Variable(0.0)
        self.m, self.s = None, None

    def apply(self, grads, variables):
        if self.m is None:
            self.m = [tf.Variable(tf.zeros_like(v)) for v in variables]
            self.s = [tf.Variable(tf.zeros_like(v)) for v in variables]
        self.t.assign_add(1.0)
        for g, v, m, s in zip(grads, variables, self.m, self.s):
            m.assign(self.b1 * m + (1 - self.b1) * g)
            s.assign(self.b2 * s + (1 - self.b2) * g * g)
            m_hat = m / (1 - self.b1 ** self.t)
            s_hat = s / (1 - self.b2 ** self.t)
            v.assign_sub(self.lr * m_hat / (tf.sqrt(s_hat) + self.eps))

tf.random.set_seed(42)
model = MLP([2, 64, 64, 3], name="spiral_net")
opt = Adam(lr=0.01)
opt.apply([tf.zeros_like(v) for v in model.trainable_variables], model.trainable_variables)  # create slots before tf.function
opt.t.assign(0.0)
print("parameters:", sum(int(tf.size(v)) for v in model.trainable_variables))

OUTPUT

parameters: 4547

💡

Why Call opt.apply With Zeros First?

Remember lesson 3 of the last module: a tf.function may not create variables after its first trace. Our Adam creates its slots on first use, so we create them once, eagerly, before the fast training step. Zero gradients leave the weights unchanged. (Keras optimizers do the same with optimizer.build().)

Section 05

Step 3 — Train Step, Accuracy and the Epoch Loop

@tf.function
def train_step(xb, yb):
    with tf.GradientTape() as tape:
        logits = model(xb)
        loss = tf.reduce_mean(
            tf.nn.sparse_softmax_cross_entropy_with_logits(labels=yb, logits=logits))
    grads = tape.gradient(loss, model.trainable_variables)
    opt.apply(grads, model.trainable_variables)
    correct = tf.reduce_sum(tf.cast(tf.argmax(logits, axis=1, output_type=tf.int32) == yb, tf.float32))
    return loss, correct

@tf.function
def eval_step(xb, yb):
    logits = model(xb)
    return tf.reduce_sum(tf.cast(tf.argmax(logits, axis=1, output_type=tf.int32) == yb, tf.float32))

EPOCHS = 60
for epoch in range(1, EPOCHS + 1):
    total_loss, total_correct, seen = 0.0, 0.0, 0
    for xb, yb in train_ds:                                   # shuffled again each epoch
        loss, correct = train_step(xb, yb)
        n = int(tf.shape(xb)[0])
        total_loss += float(loss) * n
        total_correct += float(correct)
        seen += n
    val_correct = sum(float(eval_step(xb, yb)) for xb, yb in val_ds)
    if epoch in (1, 5, 10, 20, 40, 60):
        print(f"epoch {epoch:2d}  loss {total_loss / seen:.4f}  "
              f"train acc {total_correct / seen:.3f}  val acc {val_correct / len(X_val):.3f}")

print("traces of train_step:", train_step.experimental_get_tracing_count())

OUTPUT

epoch 1 loss 0.7815 train acc 0.500 val acc 0.625 epoch 5 loss 0.0909 train acc 0.985 val acc 0.992 epoch 10 loss 0.0153 train acc 1.000 val acc 0.992 epoch 20 loss 0.0030 train acc 1.000 val acc 0.992 epoch 40 loss 0.0006 train acc 1.000 val acc 0.992 epoch 60 loss 0.0002 train acc 1.000 val acc 0.992 traces of train_step: 1

🎉

A Neural Network, Built From Nothing

Starting from random weights (about 33% accuracy — pure guessing between 3 classes), the network learns the spirals. Every part — layers, activation, loss, optimizer, loop — is code you wrote. Notice that train_step traced only once: 480 divides evenly by 32, so every batch has the same shape. With 490 samples, the last batch would hold 10 and cause a second trace (try it in the batch tool above).

Use the Trained Model

new_points = tf.constant([[0.0, 0.5], [0.4, -0.3], [-0.5, -0.1]])
probs = tf.nn.softmax(model(new_points))
for p, c in zip(probs.numpy(), tf.argmax(probs, axis=1).numpy()):
    print("probabilities", p.round(3), "-> class", c)

OUTPUT

probabilities [0. 1. 0.] -> class 1 probabilities [0.999 0. 0. ] -> class 0 probabilities [0.004 0.683 0.313] -> class 1

Section 06

Try It — Train a Network in Your Browser

🧠 Live Training Playground Interactive

This is the same MLP and training loop, written in JavaScript so it runs in your browser. Pick a dataset and settings, then press Train. The background colour shows what the network predicts at every point. Try: 1 layer with 2 units on the spiral (it fails), then 2 layers with 32 units.

data noise: 0.15 hidden layers units

optimizer learning rate batch

loss per epoch

accuracy per epoch (orange = train, blue = validation)

Section 07

Reading the Curves

What you see	What it means	What to try
Loss falls, train and val accuracy both high	Healthy training	Stop when val accuracy stops improving
Loss stuck high, accuracy near 1 / classes	Underfitting — model too small, or lr too low	More units/layers; raise lr
Loss jumps around or becomes nan	lr too high	Divide lr by 3–10
Train accuracy ≫ val accuracy	Overfitting — memorising the training set	More data, smaller model, regularisation, early stopping

Section 08

Golden Rules

🧠 Training Loops — Rules to Remember

One step = one batch. One epoch = every batch once. Steps per epoch = ⌈samples / batch size⌉.

Shuffle the training data every epoch. Never train on the validation set.

The model outputs logits; the loss is sparse_softmax_cross_entropy_with_logits; accuracy uses argmax(logits).

Wrap the train step in @tf.function, and create optimizer state before the first call.

Watch both train and validation accuracy. The gap between them tells you about overfitting.

This loop is exactly what model.fit() runs for you. Now you know every line of it.