1. Loading the data

Our Dataset is the MNIST dataset

The MNIST dataset consists of 28x28 pixel grayscale images of handwritten digits (0-9).

Dataset details:

• Training set: 60,000 images

• Test set: 10,000 images

• Image size: 28x28 pixels

• Number of classes: 10 (digits 0-9)

Data format: Each image is represented as a 2D array of pixel values, where each pixel is a grayscale intensity between 0 (black) and 255 (white).

def build_data_loader(batch_size: int) -> torch.utils.data.DataLoader:
    transform = Compose([ToTensor(), Normalize((0.5,), (0.5,))])
    train_data = MNIST(root="./data", train=True, download=True, transform=transform)
    data_loader = DataLoader(train_data, batch_size=batch_size, shuffle=True, drop_last=True)
    return data_loader

Let’s visualize a batch

fig, axs = plt.subplots(1, 10, figsize=(20, 2))

data_loader = build_data_loader(batch_size=10)

for (images, labels) in data_loader:
    
    for i, (image, label) in enumerate(zip(images, labels)):
        axs[i].imshow(image.squeeze(), cmap="gray")
        axs[i].set_title(label.item())
        axs[i].axis("off")
    break