🎓 A1Learner
🔍

PyTorch Image Batches

An image batch is a tensor with shape (N, C, H, W) that holds many images at once. N is the number of images, and C, H, and W are channels, height, and width.

Stacking a batch

Four single images stack into one batch. The new first position counts the images.

import torch

images = torch.stack([torch.zeros(1, 16, 16) for _ in range(4)])
print(images.shape)

The flatten bill

One linear layer from 256 inputs to 32 hidden units holds weights and biases. The count is 8224.

import torch

layer = torch.nn.Linear(256, 32)
print(sum(p.numel() for p in layer.parameters()))

One pass through an MLP

MLP means a stack of linear layers with activations. A batch of 4 images gives 4 rows of 3 scores.

import torch

torch.manual_seed(0)
model = torch.nn.Sequential(
    torch.nn.Flatten(),
    torch.nn.Linear(256, 32),
    torch.nn.ReLU(),
    torch.nn.Linear(32, 3),
)
batch = torch.zeros(4, 1, 16, 16)
print(model(batch).shape)

Example

import torch


def make_shape(label, gen):
    image = torch.zeros(16, 16)
    if label == 0:
        top = int(torch.randint(5, 9, (1,), generator=gen))
        image[top:top + 3, 2:14] = 1.0
    elif label == 1:
        cy = 7.5 + int(torch.randint(-1, 2, (1,), generator=gen))
        cx = 7.5 + int(torch.randint(-1, 2, (1,), generator=gen))
        yy, xx = torch.meshgrid(
            torch.arange(16), torch.arange(16), indexing="ij"
        )
        dist = torch.sqrt((yy - cy) ** 2 + (xx - cx) ** 2)
        image[(dist >= 4.5) & (dist <= 6.5)] = 1.0
    else:
        top = int(torch.randint(3, 6, (1,), generator=gen))
        left = int(torch.randint(3, 6, (1,), generator=gen))
        image[top:top + 8, left:left + 8] = 1.0
    return image


gen = torch.Generator().manual_seed(0)
labels = torch.randint(0, 3, (300,), generator=gen)
images = torch.stack([make_shape(int(label), gen) for label in labels])
images = images.reshape(300, 1, 16, 16)

train_x, train_y = images[:240], labels[:240]
test_x, test_y = images[240:], labels[240:]

torch.manual_seed(0)
model = torch.nn.Sequential(
    torch.nn.Flatten(),
    torch.nn.Linear(256, 32),
    torch.nn.ReLU(),
    torch.nn.Linear(32, 3),
)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
loss_fn = torch.nn.CrossEntropyLoss()

for step in range(150):
    loss = loss_fn(model(train_x), train_y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

params = sum(p.numel() for p in model.parameters())
with torch.no_grad():
    pred = model(test_x).argmax(dim=1)
    acc = (pred == test_y).float().mean().item()

print(f"Batch shape: {tuple(images.shape)}")
print(f"Parameters: {params}")
print(f"Test accuracy: {acc:.3f}")

Key points

  • Batch: a group of examples processed together in one pass
  • Use torch.stack to join single images into a batch
  • Flatten turns each (C, H, W) image into one long row of numbers
  • A flattened 16 by 16 image has 256 numbers, and each one needs a weight per hidden unit
  • Parameter count: the total number of weights and biases a model must learn
  • The flatten-and-MLP model is a simple baseline, and its parameter bill grows with image size
  • One batch tensor carries every image, and the model reads the whole batch in one call

In this example, the shapes generator from the earlier lesson builds 300 labelled images. The first 240 images train the model, and the last 60 test it. The model flattens each image to 256 numbers and passes them through a hidden layer of 32 units. The whole model holds 8323 parameters. Most of them sit in the first layer, where 256 inputs meet 32 units. After 150 training steps the test accuracy prints 1.000. The simple baseline solves the shapes task, and its parameter bill is the price of flattening.

Note

MLP here means a plain stack of linear layers with activations, with no image-specific parts

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

A batch tensor has shape (32, 1, 16, 16). How many images does it hold, and how big is each one?