🎓 A1Learner
🔍

PyTorch Epochs and Steps

A step is one batch through the training loop. An epoch is one full pass over the whole dataset.

Batches in one epoch

The loader length counts the batches. 100 rows in batches of 16 make 7 batches.

print(len(loader))

Steps in one epoch

Walking the loader counts one step per batch. The last batch holds the 4 leftover rows.

steps = 0
last_size = 0
for xb, yb in loader:
    steps += 1
    last_size = len(yb)
print(steps)
print(last_size)

Steps over three epochs

Each epoch repeats the same batch count. Three epochs hold three times the steps of one.

print(3 * len(loader))

Example

import torch
from torch.utils.data import TensorDataset, DataLoader

data = TensorDataset(torch.zeros(100, 1), torch.zeros(100))
loader = DataLoader(data, batch_size=16)

sizes = [len(yb) for xb, yb in loader]
print("Batch sizes in one epoch:", sizes)

total = 0
for epoch in range(1, 4):
    steps = 0
    for xb, yb in loader:
        steps += 1
    total += steps
    print(f"Epoch {epoch}: {steps} steps")
print("Total steps:", total)

Key points

  • Step: one batch through forward, loss, backward, and the update
  • Epoch: one full pass over the whole dataset
  • Steps per epoch: the number of batches in the dataset, rounded up
  • Training length is quoted in epochs, like 10 epochs
  • Epochs count passes over the data
  • Steps count single batches

In this example, the dataset holds 100 rows and the batch size is 16. One epoch splits into 6 full batches and 1 batch of 4 rows. That makes 7 steps per epoch. The loop runs 3 epochs and counts the steps in each one. The total is 21 steps.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

A run trains on 100 rows with batch size 25 for 4 epochs. How many steps is that in total?