Training Steps Calculator

A training step is one update of the model’s weights, made after the model has processed one batch of data, and an epoch is one full pass over the training set. This training steps calculator turns samples, batch size and epochs into steps per epoch, total optimizer steps, effective batch size, warmup steps and run time.

What is the difference between an epoch, a batch and a step?

The three words describe three different sizes of the same loop. A batch is a group of samples pushed through the model together. A step, also called an iteration, is the work done on one batch: forward pass, backward pass, weight update. An epoch is however many steps it takes to show the model every sample once.

The table pins each term to one run: 50,000 training samples, batch size 128, 10 epochs, a single GPU. These are the calculator’s default inputs.

Batch size is a property of the data loader, the step is a property of the optimizer, and the epoch is a property of the dataset. Mixing them up is how a schedule ends up four times too long.

How do you calculate steps per epoch and total steps?

Divide the dataset by the batch size and round up, because a partial batch still costs a step. Accumulation and extra devices each divide the count again:

• N is the number of training samples.
• B is the batch size on one device and D is the number of devices, so B_g is the global batch consumed per step.
• A is the number of gradient accumulation steps (1 means none).
• E is the number of epochs.
• The inner ceiling becomes a floor when the last incomplete batch is dropped.

The remaining outputs follow from the total:

Here w is the warmup percentage and K is the checkpoint interval in steps. Run time is the number of batches processed, \(\lceil N / B_g \rceil \times E\), multiplied by the seconds one batch takes.

Worked example: 50,000 samples, batch size 128, 10 epochs

1. Batches per epoch: \(\lceil 50000 / 128 \rceil = \lceil 390.625 \rceil = 391\). The first 390 batches are full and the last one holds 50,000 − 390 × 128 = 80 samples.
2. No accumulation, so optimizer steps per epoch are also 391 and the effective batch size is 128.
3. Total steps: \(391 \times 10 = 3910\).
4. Warmup at 10%: \(\operatorname{round}(3910 \times 0.10) = 391\) steps, exactly the first epoch.
5. Samples seen: 50,000 × 10 = 500,000.
6. Time at 0.25 seconds per batch: 3,910 × 0.25 = 977.5 s, shown as 0:16:18 (rounded to the nearest second, halves up).
7. A checkpoint every 500 steps: \(\lfloor 3910 / 500 \rfloor = 7\), the last at step 3,500.

The reverse question works the same way. A paper that trains for 10,000 steps on this setup ran for 10,000 / 391 = 25.5754 epochs, which is what the reverse mode returns.

Why are schedulers and warmup specified in steps?

A learning rate scheduler changes the rate after every optimizer update, so its clock ticks in steps. A cosine or linear decay has to know where the end is before the first update happens: PyTorch’s OneCycleLR asks for total_steps, and the Hugging Face warmup schedules ask for num_warmup_steps and num_training_steps. None of them can read the number off the data loader for you.

Get the total wrong and the shape breaks. Suppose you pass 3,910 total steps and 391 warmup steps, then switch on accumulation of 4 without recomputing. Only 980 updates happen, so the decay stops a quarter of the way down and the warmup eats 391 / 980 = 39.9% of the run instead of 10%. Once you have the right pair of numbers, the learning rate schedule calculator shows the rate at each step.

The same count drives everything else that fires on a step interval: logging, evaluation, checkpointing, and the x-axis of the loss curve you read off a cross-entropy loss plot.

How does gradient accumulation change the number of steps?

With accumulation the loop still runs a forward and backward pass on every batch, but it calls the optimizer only once every A batches, after the gradients have been summed. The data seen per epoch does not change. The number of weight updates does.

In the example, accumulating 4 batches gives an effective batch size of 128 × 4 = 512 and \(\lceil 391 / 4 \rceil = 98\) optimizer steps per epoch, so 980 over 10 epochs. Effective batch size is the number of samples whose gradients are combined into one update, and it is the figure to use when you scale the learning rate. Memory use is still set by the 128 samples loaded at once, which is the number to check in the GPU memory estimate.

Several devices behave like a bigger batch. With 1,000 samples, batch size 64 and 2 GPUs, each step consumes 128 samples, so an epoch is \(\lceil 1000 / 128 \rceil = 8\) steps instead of 16.

What does drop_last do in PyTorch’s DataLoader?

drop_last is a DataLoader argument that defaults to False. Left off, the final short batch is kept and len(dataloader) is the rounded-up quotient, 391. Set to True, the short batch is thrown away and the length is rounded down to 390, so 80 samples are skipped each epoch (a different 80 each time when shuffling is on). The behaviour is described in the official PyTorch data loading documentation.

People turn it on when a tiny last batch would upset batch normalisation statistics or when the code needs every batch to have the same shape. Over 10 epochs it moves the total from 3,910 to 3,900, enough to shift a schedule by ten steps.

Common mistakes

• Counting batches as optimizer steps while accumulating. The scheduler should be given 980, not 3,910, when A = 4.
• Using integer division for the batch count. 50000 // 128 is 390, which is only right with drop_last=True.
• Forgetting the device count. In distributed training the per-device batch is multiplied by the number of devices before dividing.
• Hard-coding warmup steps. A warmup of 391 steps is 10% of this run but about 20% of the same run at batch size 256, which has 196 steps per epoch and 1,960 in total.
• Comparing runs by step count. A step at batch 512 covers four times the data of a step at batch 128. Compare by samples seen or by epochs.
• Assuming the leftover accumulation group always updates. This calculator rounds up (98). Scripts that floor the division get 97 per epoch and 970 in total.

How do you get the number of training steps in PyTorch?

len(dataloader) is the number of batches per epoch and already respects drop_last. Everything after that is plain arithmetic.

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

dataset = TensorDataset(torch.zeros(50000, 1))
loader = DataLoader(dataset, batch_size=128, shuffle=True, drop_last=False)

epochs = 10
accumulation = 1
warmup_fraction = 0.10

batches_per_epoch = len(loader)                                   # 391
steps_per_epoch = math.ceil(batches_per_epoch / accumulation)     # 391
num_training_steps = steps_per_epoch * epochs                     # 3910
num_warmup_steps = round(num_training_steps * warmup_fraction)    # 391
print(batches_per_epoch, num_training_steps, num_warmup_steps)

print(len(DataLoader(dataset, batch_size=128, drop_last=True)))   # 390

# gradient accumulation of 4
print(math.ceil(len(loader) / 4) * epochs)   # 980  (update on the leftover group)
print(len(loader) // 4 * epochs)             # 970  (floor convention)

In Keras, model.fit(x, y, batch_size=128) on the same 50,000 samples reports 391 steps per epoch in its progress bar, and dataset.batch(128, drop_remainder=True) in tf.data is the equivalent of drop_last.

The step count says how many updates the model gets, not how large the model is. For that side of the budget, count the weights with the neural network parameter calculator, or browse the other deep learning calculators.

Frequently asked questions

How many steps are in one epoch?

Steps per epoch equal the number of training samples divided by the batch size, rounded up. With 50,000 samples and a batch size of 128 that is 391 steps, or 390 if the last incomplete batch is dropped.

Is an iteration the same as a step?

Usually yes. An iteration is one batch processed, and without gradient accumulation each batch ends in one optimizer step. With accumulation of 4, there are 4 iterations for every optimizer step.

How do I convert steps to epochs?

Divide the number of steps by the steps per epoch. With 391 steps per epoch, 10,000 steps are 10,000 / 391 = 25.5754 epochs, so 26 whole epochs are needed to reach that many steps.

Does a larger batch size mean fewer training steps?

Yes. The model sees the same data per epoch in fewer, larger updates. Going from batch size 128 to 256 on 50,000 samples cuts an epoch from 391 steps to 196.

Related calculators

Tools that fit around a training plan:

• Warmup and decay schedules, step by step
• GPU memory needed for a model and batch
• Parameter count of a network, layer by layer
• Output size of a Conv2d layer
• All deep learning calculators