Load the MNIST data

From here on we're building one real program. By the end of Chapter 6 the code from these four chapters stacks into a single train.py that you run once. This chapter gets the data.

Our project is the classic first task in deep learning: recognise handwritten digits. The dataset is MNIST — 70,000 little 28×28 grayscale images of the digits 0–9, each labelled with the number it shows. It's the "hello world" of neural networks because it's small, real, and a simple model can nail it.

Start your train.py

Create a file called train.py. We'll grow it chapter by chapter. Start with the imports and the device pick from Chapter 1:

import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

device = "cuda" if torch.cuda.is_available() else "cpu"

Download the digits

torchvision can fetch MNIST for us. The transform converts each image from its raw form into a tensor with pixel values scaled to 0.0–1.0 — exactly what the network wants. We grab two separate splits: train (60,000 images) to learn from, and test (10,000 held-out images) to grade ourselves on later:

transform = transforms.ToTensor()

train_data = datasets.MNIST(
  root="data", train=True, download=True, transform=transform
)
test_data = datasets.MNIST(
  root="data", train=False, download=True, transform=transform
)

print("training images:", len(train_data))   # 60000
print("test images:", len(test_data))         # 10000

The first run downloads ~10 MB into a data/ folder; after that it loads instantly. Each item is an (image, label) pair — the image a 1×28×28 tensor, the label an integer 0–9:

image, label = train_data[0]
print(image.shape)   # torch.Size([1, 28, 28])
print(label)         # 5

Serve it in batches with a DataLoader

We don't feed the network one image at a time, and we don't feed it all 60,000 at once. We feed batches — small groups — which trains faster and more stably. A DataLoader handles the batching and shuffling for us. Add this to train.py:

train_loader = DataLoader(train_data, batch_size=64, shuffle=True)
test_loader = DataLoader(test_data, batch_size=1000)

Looping over train_loader now yields batches of 64 images and their 64 labels at a time. We shuffle the training data so the network doesn't learn anything from the order; the test data needs no shuffling. That's the entire data pipeline — next we build the network that will learn from it.