🎓 A1Learner
🔍
0 of 87 lessons completed

PyTorch Batch Size Effects

Batch size is the number of rows a model processes in one step. The size changes how noisy each step is and how many rows one step costs.

Batches per epoch

An epoch is one full pass over the data. A dataset of 256 rows splits into far fewer batches at size 128 than at size 8.

import torch

rows = 256
for batch_size in [8, 32, 128]:
    count = (rows + batch_size - 1) // batch_size
    print(f"Batch size {batch_size}: {count} batches")

Noise in small batches

One untrained model scores five random batches of each size. The smallest batches swing widest, from 0.5312 to 0.9378.

import torch

torch.manual_seed(0)
centers = torch.tensor([[-2.0, -2.0], [2.0, 2.0]])
parts = []
labels = []
for label, center in enumerate(centers):
    parts.append(center + torch.randn(128, 2) * 0.8)
    labels.extend([label] * 128)
x_all = torch.cat(parts)
y_all = torch.tensor(labels)
loss_fn = torch.nn.CrossEntropyLoss()
torch.manual_seed(1)
model = torch.nn.Linear(2, 2)
for batch_size in [8, 32, 128]:
    generator = torch.Generator().manual_seed(5)
    losses = []
    for i in range(5):
        idx = torch.randint(0, len(x_all), (batch_size,), generator=generator)
        with torch.no_grad():
            losses.append(loss_fn(model(x_all[idx]), y_all[idx]).item())
    print(f"Batch size {batch_size}: min {min(losses):.4f}, max {max(losses):.4f}")

Rows seen per step

One step at size 128 reads sixteen times more rows than one step at size 8. Step counts alone miss that difference.

steps = 46
for batch_size in [8, 128]:
    print(f"Batch size {batch_size}: {steps * batch_size} rows in {steps} steps")

Example

import torch

torch.manual_seed(0)
centers = torch.tensor([[-2.0, -2.0], [2.0, 2.0]])
parts = []
labels = []
for label, center in enumerate(centers):
    parts.append(center + torch.randn(128, 2) * 0.8)
    labels.extend([label] * 128)
x_all = torch.cat(parts)
y_all = torch.tensor(labels)
loss_fn = torch.nn.CrossEntropyLoss()
target = 0.05


def steps_to_target(batch_size):
    torch.manual_seed(1)
    model = torch.nn.Linear(2, 2)
    optimizer = torch.optim.SGD(model.parameters(), lr=0.05)
    steps = 0
    epoch = 0
    while steps < 2000:
        generator = torch.Generator().manual_seed(100 + epoch)
        order = torch.randperm(len(x_all), generator=generator)
        for start in range(0, len(x_all), batch_size):
            idx = order[start:start + batch_size]
            optimizer.zero_grad()
            loss = loss_fn(model(x_all[idx]), y_all[idx])
            loss.backward()
            optimizer.step()
            steps += 1
            with torch.no_grad():
                full = loss_fn(model(x_all), y_all).item()
            if full < target:
                return steps
        epoch += 1
    return steps


for batch_size in [8, 32, 128]:
    steps = steps_to_target(batch_size)
    rows_seen = steps * batch_size
    print(
        f"Batch size {batch_size}: {steps} steps, "
        f"{rows_seen} rows seen to loss below {target}"
    )

Key points

  • Batch size: the number of rows in one training step
  • Small batches give noisy losses, and the spread shrinks as batches grow
  • Use steps times batch size to count the rows a run has seen
  • On this task, all three sizes needed about the same number of steps
  • The batch-128 run saw about sixteen times more rows than the batch-8 run
  • Batch size trades step noise against rows per step, and neither extreme wins by default

In this example, one classifier trains on 256 blob points at three batch sizes. Each run reseeds the model, so all three starts match. Each run shuffles with a seeded generator, so the orders repeat. Training stops when the full-data loss first falls below 0.05. Batch size 8 needs 46 steps and sees 368 rows. Batch size 32 needs 47 steps and sees 1504 rows. Batch size 128 needs 46 steps and sees 5888 rows. The step counts match closely, and the row counts do not.

Note

Steps are not rows, and comparing runs by steps alone hides the data each step consumed

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

Two runs both took 46 steps, one at batch size 8 and one at batch size 128. What is the big difference between them?