PyTorch Mini-Batches by Hand

A mini-batch is a small slice of a dataset used for one training step. Cutting batches by hand means slicing with index ranges inside the training loop.

Slicing by hand, line after line

Three batches mean three slicing lines. A longer dataset would need even more lines.

batch_one = x[0:4]
batch_two = x[4:8]
batch_three = x[8:12]
print(batch_one.shape)
print(batch_two.shape)
print(batch_three.shape)

Index ranges in a loop

A range can walk the batch starts instead. Each start steps up by the batch size.

for start in range(0, 12, 4):
    print(start, start + 4)

One hand-cut training step

One slice feeds one normal training step. The step itself does not change.

optimizer.zero_grad()
loss = loss_fn(model(x[0:4]), y[0:4])
loss.backward()
optimizer.step()
print(f"Loss: {loss.item():.4f}")

Example

import torch
import torch.nn as nn

torch.manual_seed(0)
x = torch.arange(12, dtype=torch.float32).reshape(-1, 1) / 10.0
y = 2.0 * x + 1.0

model = nn.Linear(1, 1)
loss_fn = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.5)

for epoch in range(1, 11):
    total = 0.0
    for start in range(0, 12, 4):
        xb = x[start:start + 4]
        yb = y[start:start + 4]
        optimizer.zero_grad()
        loss = loss_fn(model(xb), yb)
        loss.backward()
        optimizer.step()
        total += loss.item()
    print(f"Epoch {epoch}: loss {total / 3:.4f}")

print(f"Weight: {model.weight.item():.2f}")
print(f"Bias: {model.bias.item():.2f}")

Key points

Mini-batch: a small slice of rows used for one training step
Use range(0, rows, batch_size) to walk the start of each slice
Use x[start:start + batch_size] to cut one batch of rows
One epoch of batches still uses every row once
Hand-cut batches are plain slices, cut fresh in every epoch

In this example, x holds 12 rows and y holds the matching targets. The inner loop cuts batches of 4 rows with index ranges. Each batch runs one full training step. The printed loss is the average of the three batch losses. After 10 epochs the weight and bias sit near the true values 2 and 1.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

import torch
import torch.nn as nn

torch.manual_seed(0)
x = torch.arange(12, dtype=torch.float32).reshape(-1, 1) / 10.0
y = 2.0 * x + 1.0

model = nn.Linear(1, 1)
loss_fn = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.5)

for epoch in range(1, 11):
    total = 0.0
    for start in range(0, 12, 4):
        xb = x[start:start + 4]
        yb = y[start:start + 4]
        optimizer.zero_grad()
        loss = loss_fn(model(xb), yb)
        loss.backward()
        optimizer.step()
        total += loss.item()
    print(f"Epoch {epoch}: loss {total / 3:.4f}")

print(f"Weight: {model.weight.item():.2f}")
print(f"Bias: {model.bias.item():.2f}")

Exercise

A dataset has 12 rows and the batch size is 4. How many batches make one epoch?

3 4 12 48