Batch Size and Epochs

Batch size is the number of samples processed together. An epoch is one complete pass through the training dataset.

In a standard training loop, the model calculates a loss for each batch and performs one optimiser update. This update is one training step.

Counting steps

Suppose the training set contains 10 samples and the batch size is 4. Keep the final incomplete batch:

Batch	Samples	Updates
1	4	1
2	4	1
3	2	1

Therefore, one epoch has 3 updates. Two epochs have 6 updates, and each sample is used twice.

For $N$ samples and batch size $B$:

$$\text{steps per epoch}=\lceil \frac{N}{B}\rceil$$

The ceiling symbol means round up. This count assumes one update per batch, with the final batch retained.

Choosing the values

Effects to compare:

• A larger batch needs more memory and uses more samples in each gradient estimate
• A smaller batch gives more updates per epoch and usually more variation between gradient estimates
• Training speed depends on the hardware, model, and data loading
• Validation results determine which batch size works better for the task
• More epochs give the model more training opportunities, but can increase overfitting

Early Stopping can stop training when validation performance stops improving.

Code

PyTorch

import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset
 
X = torch.arange(10, dtype=torch.float32).reshape(-1, 1) / 10
y = 2 * X
loader = DataLoader(
    TensorDataset(X, y), batch_size=4, shuffle=True, drop_last=False
)
model = nn.Linear(1, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
criterion = nn.MSELoss()
 
for epoch in range(2):
    sizes = []
    for data, target in loader:
        optimizer.zero_grad()
        loss = criterion(model(data), target)
        loss.backward()
        optimizer.step()
        sizes.append(len(data))
    print(epoch + 1, sizes)

1 [4, 4, 2]
2 [4, 4, 2]

Read the loops:

• reshape(-1, 1) makes 10 rows with one feature per row
• TensorDataset pairs each input row with its target row
• DataLoader supplies the batches and shuffles the sample order each epoch
• drop_last=False retains the last 2 samples
• The outer loop selects the epoch, while the inner loop selects the batch
• Each call to optimizer.step() performs one update

With drop_last=True, each epoch would use 2 full batches and omit the final 2 samples. See the DataLoader API.

TensorFlow and Keras

import numpy as np
import tensorflow as tf
 
X = np.arange(10, dtype="float32").reshape(-1, 1) / 10
y = 2 * X
model = tf.keras.Sequential([
    tf.keras.Input(shape=(1,)),
    tf.keras.layers.Dense(1),
])
model.compile(
    optimizer=tf.keras.optimizers.SGD(learning_rate=0.01), loss="mse"
)
history = model.fit(
    X, y, batch_size=4, epochs=2, shuffle=True, verbose=0
)
print(int(model.optimizer.iterations.numpy()))

6

Read the equivalent controls:

• fit() creates the batches and performs both loops
• batch_size=4 gives 3 batches per epoch for these arrays
• epochs=2 requests 2 passes through the data
• history.history["loss"] contains one training loss value per epoch
• optimizer.iterations counts the 6 updates made by this new optimiser