Explore PyTorch datasets and models#

The introductory workshop used handwritten digits and a custom CNN. This notebook asks two broader questions:

  1. What other datasets can we explore with the same PyTorch workflow?

  2. How can we use PyTorch’s popular, established models and their pretrained weights for our own tasks?

We will adapt a custom model to a new dataset, then use ImageNet-pretrained ResNet to explore transfer learning.

1. The dataset landscape#

PyTorch’s DataLoader works with any object following the Dataset interface: it provides a length and returns one sample at a time. Domain libraries offer ready-made datasets, and you can write your own Dataset for your research data.

Data or task

A few examples

Where to explore

Small grayscale images

MNIST, FashionMNIST

Torchvision image datasets

Color image classification

CIFAR-10, CIFAR-100

CIFAR-10 example

Locate or segment objects

COCO Detection, VOC Segmentation

Torchvision detection and segmentation datasets

Speech or audio

YESNO, Speech Commands

TorchAudio datasets

Your own text, tables, signals, or images

A custom Dataset returning input and target pairs

PyTorch Dataset and DataLoader

These are examples, not a complete list.

Explore: Find a dataset relevant to your field. What would one input look like? What would the target be? Would you need classification, regression, detection, or another task?

A glimpse of another image dataset#

CIFAR-10 uses 32 × 32 RGB images across ten object classes. Its loading pattern resembles MNIST:

# Example only; do not run this for today's challenge.
# another_dataset = torchvision.datasets.CIFAR10(
#     root="./data", train=True, download=True,
#     transform=transforms.ToTensor()
# )

The input shape and target labels would change, but Dataset, DataLoader, model, loss, and evaluation remain the main pieces of the workflow.

2. Our hands-on challenge: EMNIST#

EMNIST extends handwritten-character recognition beyond the ten MNIST digits. We will use the balanced split, which has 47 classes representing digits and letters. Its images are still grayscale and 28 × 28 pixels, so a CNN from the earlier challenge is a useful starting point. The final layer must now output one logit per EMNIST class.

Setup and experiment controls#

Start with a subset so an experiment finishes quickly. Once everything works, increase the sample limits or set them to None for the full split. Change one setting at a time and record what happens.

import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader, Subset

SEED = 355
torch.manual_seed(SEED)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print("Device:", device)
if not torch.cuda.is_available():
    print("A CUDA GPU was not detected; training will be slower.")

# CHANGE ME: experiment with these after your first run.
BATCH_SIZE = 64
learning_rate = 0.001
num_epochs = 2

NUM_WORKERS = 8
DATA_ROOT = "./data"

Load the dataset#

For a fair comparison, the sample indices below are chosen once with a fixed seed. The test loader is never used to update model parameters.

transform = transforms.Compose([
    transforms.ToTensor()
])

trainset = torchvision.datasets.EMNIST(
    root=DATA_ROOT,
    split="balanced",
    train=True,
    download=True,
    transform=transform,
)
testset = torchvision.datasets.EMNIST(
    root=DATA_ROOT,
    split="balanced",
    train=False,
    download=True,
    transform=transform,
)

# Get the class count from the dataset, rather than assuming MNIST's ten.
classes = trainset.classes
num_classes = len(classes)

trainloader = DataLoader(
    trainset, batch_size=BATCH_SIZE, shuffle=True, num_workers=NUM_WORKERS
)
testloader = DataLoader(
    testset, batch_size=BATCH_SIZE, shuffle=False, num_workers=NUM_WORKERS
)

print("Training examples in this run:", len(trainset))
print("Test examples in this run:", len(testset))
print("Number of classes:", num_classes)
print("First ten class names:", classes[:10])
assert num_classes == 47

Inspect a batch#

Before you run the cell, predict the shape of the image and label batches. Which dimensions are the batch size, channels, height, and width? If some handwritten characters look unusual, look at several samples before judging the model.

images, labels = next(iter(trainloader))
print("Image batch shape:", images.shape)
print("Label batch shape:", labels.shape)
print("First labels:", [classes[label] for label in labels[:10].tolist()])

grid = torchvision.utils.make_grid(images[:16], nrow=8)
plt.figure(figsize=(12, 4))
plt.imshow(np.transpose(grid.numpy(), (1, 2, 0)), cmap="gray")
plt.axis("off")
plt.show()

3. Adapt a custom model#

You may reuse the two-convolution model from the end of the first notebook. The starter below follows that design. Compare the old task with EMNIST:

  • The images are still 1 × 28 × 28, so the first convolution can stay the same.

  • Two unpadded 3 × 3 convolutions still change 28 → 26 → 24.

  • The output layer must change from 10 logits to num_classes logits.

Try your own architecture if you prefer. A shape mismatch at the first linear layer means you need to recalculate its input size.

class MyModel2(nn.Module):
    def __init__(self, num_classes):
        super(MyModel2, self).__init__()

        # CHANGE ME: replace these layers with your model from the
        # earlier challenge, or try a new architecture.
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=16, kernel_size=3)
        self.conv2 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3)
        self.d1 = nn.Linear(32 * 24 * 24, 128)
        self.d2 = nn.Linear(128, num_classes)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = x.flatten(start_dim=1)
        x = F.relu(self.d1(x))
        return self.d2(x)  # Return logits for CrossEntropyLoss.

Check the model before training#

The output must have one row per image and one column per class. Rerun this cell for every new architecture so each experiment starts with fresh weights.

torch.manual_seed(SEED)
model = MyModel2(num_classes=num_classes).to(device)

with torch.no_grad():
    sample_logits = model(images.to(device))

print(model)
print("Input shape:", images.shape)
print("Output shape:", sample_logits.shape)
print("Trainable parameters:", f"{sum(p.numel() for p in model.parameters()):,}")
assert sample_logits.shape == (images.size(0), num_classes)

4. Train and evaluate#

This is the same forward pass → loss → gradients → optimizer update pattern as before. The test set is evaluated after each epoch so you can compare training performance with performance on unseen characters.

learning-rate schedulers. A scheduler adjusts the learning rate during training. Here, StepLR halves it after each epoch (step_size=1, gamma=0.5), so the first two epochs use 0.001 and 0.0005. Smaller updates can help refine the model as training progresses. Call scheduler.step() once at the end of each epoch, after the optimizer updates.

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.5)
history = []

for epoch in range(num_epochs):
    train_running_loss = 0.0
    train_correct = 0
    train_examples = 0

    model.train()
    for images, labels in trainloader:
        images = images.to(device)
        labels = labels.to(device)

        outs = model(images)
        loss = criterion(outs, labels)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        train_running_loss += loss.detach().item() * labels.size(0)
        train_correct += (outs.argmax(dim=1) == labels).sum().item()
        train_examples += labels.size(0)

    test_correct = 0
    test_examples = 0
    model.eval()
    with torch.no_grad():
        for images, labels in testloader:
            images = images.to(device)
            labels = labels.to(device)
            outputs = model(images)
            test_correct += (outputs.argmax(dim=1) == labels).sum().item()
            test_examples += labels.size(0)

    train_accuracy = 100 * train_correct / train_examples
    test_accuracy = 100 * test_correct / test_examples
    history.append((train_accuracy, test_accuracy))
    print(
        f"Epoch {epoch + 1}/{num_epochs} | "
        f"Loss: {train_running_loss / train_examples:.4f} | "
        f"Train accuracy: {train_accuracy:.2f}% | "
        f"Test accuracy: {test_accuracy:.2f}% | "
        f"Learning rate: {scheduler.get_last_lr()[0]:.6f}"
    )
    scheduler.step()  # Set the learning rate for the next epoch.

Investigate predictions#

Look beyond the overall accuracy. Which characters does the model confuse? Could a human also confuse those handwritten shapes?

images, labels = next(iter(testloader))
model.eval()
with torch.no_grad():
    outputs = model(images.to(device))
    probabilities = F.softmax(outputs, dim=1).cpu()

fig, axes = plt.subplots(2, 8, figsize=(14, 4))
for i in range(8):
    prediction = probabilities[i].argmax().item()
    axes[0, i].imshow(images[i].squeeze(), cmap="gray")
    axes[0, i].set_title(f"true: {classes[labels[i].item()]}")
    axes[0, i].axis("off")

    axes[1, i].bar(np.arange(num_classes), probabilities[i].numpy())
    axes[1, i].set_ylim(0, 1)
    axes[1, i].set_xticks([])
    axes[1, i].set_title(f"pred: {classes[prediction]}")

plt.tight_layout()
plt.show()

Your EMNIST experiment#

Change one thing, then rerun the model creation, training, and evaluation cells. Ideas: add pooling, add dropout, change the number of channels, change the hidden layer size, or train for another epoch. If you change batch size, rerun the data-loading cells too.

Run

What changed?

Train accuracy

Test accuracy

Baseline

Starter model

Experiment 1

Experiment 2

Did a change help both training and test accuracy? If it helped only training accuracy, what might that tell you?

5. Explore existing models and transfer learning#

PyTorch offers popular model architectures with established performance on image recognition benchmarks through Torchvision’s model library. These include ResNet, MobileNet, EfficientNet, and Vision Transformer. We can use these models as starting points for our own tasks. Many have downloadable weights, meaning parameter values learned during earlier training.

Pretraining trains a model on an earlier task, often ImageNet image classification. Transfer learning adapts that learned model to a new task. Two common approaches are:

  1. Fixed feature extractor: freeze the pretrained layers, replace the final classifier, and train only the new classifier.

  2. Fine-tuning: start from pretrained weights and update some or all layers on the new task.

The first is cheaper to try. The second can adapt more deeply but needs more compute and careful tuning. PyTorch transfer learning tutorial

Browse available model architectures#

Torchvision includes more than one ResNet size, as well as other model families. This quick inventory does not download weights or train anything. Which architecture would you want to investigate for a different image project? Torchvision model listing

from torchvision.models import list_models

# List a few available ResNet architecture names; weights are separate.
resnet_names = [name for name in list_models() if name.startswith("resnet")]
print("Available ResNet architectures:", resnet_names[:8])

Adapt ResNet-18 to EMNIST#

ResNet-18’s ImageNet weights came from color natural images and a 1,000-class output layer. EMNIST contains grayscale 28 × 28 characters and 47 classes. We therefore need to:

  • convert each grayscale image to three channels;

  • apply the preprocessing paired with the pretrained weights, including resizing and normalization;

  • replace the final 1,000-class layer with a 47-class layer.

We will reuse the pretrained feature extractor and train a new classifier for our character labels. ResNet-18 weights and preprocessing

ResNet-18 experiment#

Run the following cells to load the pretrained weights and adapt ResNet-18. The weights download on the first run. We use a small sample because the ImageNet preprocessing enlarges each 28 × 28 image.

from torchvision.models import ResNet18_Weights, resnet18

weights = ResNet18_Weights.IMAGENET1K_V1
transfer_transform = transforms.Compose([
    transforms.Grayscale(num_output_channels=3),
    weights.transforms(),
])

# Reuse the already downloaded EMNIST files with new preprocessing.
transfer_trainset = torchvision.datasets.EMNIST(
    root=DATA_ROOT, split="balanced", train=True,
    download=False, transform=transfer_transform
)
transfer_testset = torchvision.datasets.EMNIST(
    root=DATA_ROOT, split="balanced", train=False,
    download=False, transform=transfer_transform
)
transfer_trainloader = DataLoader(
    transfer_trainset,
    batch_size=16, shuffle=True, num_workers=NUM_WORKERS
)
transfer_testloader = DataLoader(
    transfer_testset,
    batch_size=16, shuffle=False, num_workers=NUM_WORKERS
)

# Download ImageNet-pretrained ResNet-18 weights if not already cached.
transfer_model = resnet18(weights=weights)
for parameter in transfer_model.parameters():
    parameter.requires_grad = False

# Replace the ImageNet classifier with a trainable EMNIST classifier.
transfer_model.fc = nn.Linear(transfer_model.fc.in_features, num_classes)
transfer_model = transfer_model.to(device)
transfer_model.eval()  # Keep pretrained batch-normalization statistics fixed.

transfer_images, transfer_labels = next(iter(transfer_trainloader))
with torch.no_grad():
    transfer_logits = transfer_model(transfer_images.to(device))
print("Preprocessed input shape:", transfer_images.shape)
print("New output shape:", transfer_logits.shape)
print("Trainable parameters:",
      f"{sum(p.numel() for p in transfer_model.parameters() if p.requires_grad):,}")
assert transfer_logits.shape == (transfer_images.size(0), num_classes)

Train only the new classifier#

The pretrained backbone remains frozen. We keep it in evaluation mode so its stored batch-normalization statistics do not change while training the new final layer. Compare the result with your custom CNN, remembering that this run uses fewer examples.

# CHANGE ME: experiment with these hyperparameters.
learning_rate = 0.001
num_epochs = 5
lr_step_size = 1  # Reduce the learning rate every epoch.
lr_gamma = 0.5    # Multiply the learning rate by this factor.

transfer_criterion = nn.CrossEntropyLoss()
transfer_optimizer = torch.optim.Adam(
    transfer_model.fc.parameters(), lr=learning_rate
)
transfer_scheduler = torch.optim.lr_scheduler.StepLR(
    transfer_optimizer,
    step_size=lr_step_size,
    gamma=lr_gamma,
)

for epoch in range(num_epochs):
    train_losses = []
    train_correct = 0
    train_examples = 0

    # Keep the frozen backbone's batch-normalization statistics fixed.
    transfer_model.eval()
    transfer_model.fc.train()

    for transfer_images, transfer_labels in transfer_trainloader:
        transfer_images = transfer_images.to(device)
        transfer_labels = transfer_labels.to(device)

        transfer_optimizer.zero_grad()
        transfer_outputs = transfer_model(transfer_images)
        transfer_loss = transfer_criterion(transfer_outputs, transfer_labels)
        transfer_loss.backward()
        transfer_optimizer.step()

        # Save each batch's loss so we can calculate the mean.
        train_losses.append(transfer_loss.item())
        train_correct += (
            transfer_outputs.argmax(dim=1) == transfer_labels
        ).sum().item()
        train_examples += transfer_labels.size(0)

    test_correct = 0
    test_examples = 0
    transfer_model.eval()

    # Evaluate without computing gradients or updating parameters.
    with torch.no_grad():
        for transfer_images, transfer_labels in transfer_testloader:
            transfer_images = transfer_images.to(device)
            transfer_labels = transfer_labels.to(device)

            predictions = transfer_model(transfer_images).argmax(dim=1)
            test_correct += (predictions == transfer_labels).sum().item()
            test_examples += transfer_labels.size(0)

    train_accuracy = 100 * train_correct / train_examples
    test_accuracy = 100 * test_correct / test_examples
    print(
        f"Epoch {epoch + 1}/{num_epochs} | "
        f"Loss: {np.mean(train_losses):.4f} | "
        f"Train accuracy: {train_accuracy:.2f}% | "
        f"Test accuracy: {test_accuracy:.2f}% | "
        f"Learning rate: {transfer_scheduler.get_last_lr()[0]:.6f}"
    )
    transfer_scheduler.step()  # Set the learning rate for the next epoch.

Final takeaways#

  • Datasets. PyTorch offers datasets for many tasks beyond image classification. We explored these options and used EMNIST to extend handwritten-digit recognition to digits and letters.

  • Hyperparameter tuning. Experiment with settings such as learning rate, batch size, and number of epochs.

  • Models and transfer learning. Torchvision provides popular, established models such as ResNet. We reused ImageNet-pretrained weights and trained a new classifier for EMNIST, showing how an existing model can be adapted and fine-tuned to a new task.