🎓 A1Learner
🔍

PyTorch The Shapes Dataset

A dataset class is a class that serves one data row at a time through __len__ and __getitem__. The shapes dataset draws bars, rings, and squares as tensors, so the whole dataset is made by code.

A bar from two comparisons

Rows and columns start as index grids. A comparison on each grid gives a mask of True and False cells.

rows = torch.arange(16).view(16, 1)
cols = torch.arange(16).view(1, 16)
mask = (rows >= 7) & (rows <= 9) & (cols >= 0)
print(mask.shape)
print(int(mask.sum()))

A ring from distances

Distance from the centre decides the ring. Cells at distance 4.5 to 6.5 from the centre belong to the ring.

rows = torch.arange(16).view(16, 1).float()
cols = torch.arange(16).view(1, 16).float()
dist = torch.sqrt((rows - 7.5) ** 2 + (cols - 7.5) ** 2)
ring = (dist >= 4.5) & (dist <= 6.5)
print(int(ring.sum()))

The Dataset contract

__getitem__ returns one image and its label. The image carries one channel, so its shape is (1, 16, 16).

data = ShapesDataset(300, seed=42)
image, label = data[0]
print(len(data))
print(image.shape, label)

The same seed, the same dataset

A fixed seed makes the whole dataset repeatable. Two builds with seed 42 hold the same images.

again = ShapesDataset(300, seed=42)
print(torch.equal(data.images, again.images))

Example

import torch
from torch.utils.data import Dataset


def draw_shape(label, g):
    rows = torch.arange(16).view(16, 1).float()
    cols = torch.arange(16).view(1, 16).float()
    if label == 0:
        centre = 8 + int(torch.randint(-2, 3, (1,), generator=g))
        mask = (rows >= centre - 1) & (rows <= centre + 1) & (cols >= 0)
    elif label == 1:
        cr = 7.5 + int(torch.randint(-1, 2, (1,), generator=g))
        cc = 7.5 + int(torch.randint(-1, 2, (1,), generator=g))
        dist = torch.sqrt((rows - cr) ** 2 + (cols - cc) ** 2)
        mask = (dist >= 4.5) & (dist <= 6.5)
    else:
        top = int(torch.randint(2, 7, (1,), generator=g))
        left = int(torch.randint(2, 7, (1,), generator=g))
        mask = (rows >= top) & (rows <= top + 7) & (cols >= left) & (cols <= left + 7)
    return mask.float()


class ShapesDataset(Dataset):
    def __init__(self, count, seed):
        g = torch.Generator().manual_seed(seed)
        labels = torch.arange(count) % 3
        self.labels = labels[torch.randperm(count, generator=g)]
        images = [draw_shape(int(label), g) for label in self.labels]
        self.images = torch.stack(images).unsqueeze(1)

    def __len__(self):
        return len(self.labels)

    def __getitem__(self, index):
        return self.images[index], self.labels[index]


data = ShapesDataset(300, seed=42)

split = torch.randperm(300, generator=torch.Generator().manual_seed(7))
train_idx = split[:240]
val_idx = split[240:]

names = ["bar", "ring", "square"]
print("Dataset size:", len(data))
print("Train size:", len(train_idx))
print("Validation size:", len(val_idx))
for label in range(3):
    count = int((data.labels[train_idx] == label).sum())
    print(f"Train class {label} ({names[label]}):", count)

for label in range(3):
    first = int((data.labels == label).nonzero()[0])
    image, _ = data[first]
    print(f"One {names[label]} (class {label}):")
    for row in image[0].tolist():
        print("".join("#" if value == 1.0 else "." for value in row))

Key points

  • Dataset class: a class whose __getitem__ returns one row and whose __len__ counts the rows
  • Labels: 0 for a bar, 1 for a ring, 2 for a square
  • Use torch.arange grids and comparisons to build a mask for each shape
  • Use torch.sqrt on row and column distances to draw a ring
  • Use one seeded torch.Generator for every random choice, so the dataset repeats exactly
  • Use a seeded randperm on the indices to split off a validation part

In this example, draw_shape builds one image from masks and distances. The ShapesDataset class draws 300 images with a generator seeded at 42. The labels start as 100 of each class and are shuffled by the same generator. A second seeded randperm splits the indices into 240 train rows and 60 validation rows. The train counts print at 78 bars, 79 rings, and 83 squares. One sample per class then prints as a shade grid. The bar fills three full rows. The ring is a hollow circle. The square is a filled 8 by 8 block.

Note

The images hold only 0.0 and 1.0 values, so a shade grid prints each image exactly.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

What does __getitem__ return for the shapes dataset?