Convolution Neural Networks

Shan-Hung Wu & DataLab
Fall 2026

In this lab, we introduce two datasets, MNIST and CIFAR, then we will talk about how to implement CNN models for these two datasets using PyTorch. The major difference between mnist and cifar is their size. Due to the limit of memory size and time issue, we offer a guide to illustrate a typical data loading pipeline in PyTorch. Let's dive into PyTorch!

import warnings
import os
warnings.filterwarnings("ignore")

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

import matplotlib as mpl
import pickle
import csv
import itertools
from collections import defaultdict
import time
import pandas as pd
import math
from tqdm import tqdm

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print('Using device:', device)

Using device: cuda

import urllib.request
if not os.path.exists("lab11_1_lib_torch.py"):
    urllib.request.urlretrieve("https://nthu-datalab.github.io/ml/labs/11-1_CNN/lab11_1_lib_torch.py", "lab11_1_lib_torch.py")

from lab11_1_lib_torch import draw_timeline

MNIST

We start from a simple dataset. MNIST is a simple computer vision dataset. It consists of images of handwritten digits like:

It also includes labels for each image, telling us which digit it is. For example, the labels for the above images are 5, 0, 4, and 1. Each image is 28 pixels by 28 pixels. We can interpret this as a big array of numbers:

The MNIST data is hosted on Yann LeCun's website. We can directly import the MNIST dataset via torchvision.datasets.

# Download and prepare the MNIST dataset
mnist_train = torchvision.datasets.MNIST(root='./data', train=True, download=True)
mnist_test = torchvision.datasets.MNIST(root='./data', train=False, download=True)

# Normalize pixel values to be between 0 and 1
train_image = mnist_train.data.numpy().astype('float32') / 255.0
train_label = mnist_train.targets.numpy()
test_image = mnist_test.data.numpy().astype('float32') / 255.0
test_label = mnist_test.targets.numpy()

print('shape of train_image:', train_image.shape)
print('shape of train_label:', train_label.shape)

shape of train_image: (60000, 28, 28)
shape of train_label: (60000,)

Softmax Regression on MNIST

Before jumping to Convolutional Neural Network model, we're going to start with a very simple model with a single layer and softmax regression.

We know that every image in MNIST is a handwritten digit between zero and nine. So there are only ten possible digits that a given image can be. We want to give the probability of the input image for being each digit. That is, input an image, the model outputs a ten-dimension vector.

This is a classic case where a softmax regression is a natural, simple model. If you want to assign probabilities to an object being one of several different things, softmax is the thing to do.

# flatting the training data for dense layers
train_image_1 = train_image.reshape((60000, -1))
test_image_1 = test_image.reshape((10000, -1))
print(train_image_1.shape)
print(test_image_1.shape)

(60000, 784)
(10000, 784)

# Softmax regression == a single Linear layer (CrossEntropyLoss applies the
# softmax internally, so the model itself just outputs logits)
model_1 = nn.Linear(784, 10).to(device)
print(model_1)

Linear(in_features=784, out_features=10, bias=True)

# compile the model and train it for 5 epochs
def train_sklearn_style(model, X, y, epochs, batch_size=128, lr=1e-3):
    """Small helper that mimics keras' `model.fit(X, y, epochs=...)` one-liner."""
    ds = torch.utils.data.TensorDataset(torch.from_numpy(X), torch.from_numpy(y).long())
    loader = DataLoader(ds, batch_size=batch_size, shuffle=True)
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    criterion = nn.CrossEntropyLoss()

    model.train()
    for epoch in range(epochs):
        total_loss, correct, n = 0.0, 0, 0
        for xb, yb in loader:
            xb, yb = xb.to(device), yb.to(device)
            optimizer.zero_grad()
            logits = model(xb)
            loss = criterion(logits, yb)
            loss.backward()
            optimizer.step()

            total_loss += loss.item() * xb.size(0)
            correct += (logits.argmax(1) == yb).sum().item()
            n += xb.size(0)
        print(f'Epoch {epoch+1}/{epochs} - loss: {total_loss/n:.4f} - accuracy: {correct/n:.4f}')

train_sklearn_style(model_1, train_image_1, train_label, epochs=5)

Epoch 1/5 - loss: 0.6617 - accuracy: 0.8499
Epoch 2/5 - loss: 0.3594 - accuracy: 0.9029
Epoch 3/5 - loss: 0.3178 - accuracy: 0.9133
Epoch 4/5 - loss: 0.2985 - accuracy: 0.9172
Epoch 5/5 - loss: 0.2868 - accuracy: 0.9201

@torch.no_grad()
def evaluate(model, X, y, batch_size=256):
    model.eval()
    ds = torch.utils.data.TensorDataset(torch.from_numpy(X), torch.from_numpy(y).long())
    loader = DataLoader(ds, batch_size=batch_size)
    correct, n = 0, 0
    for xb, yb in loader:
        xb, yb = xb.to(device), yb.to(device)
        correct += (model(xb).argmax(1) == yb).sum().item()
        n += xb.size(0)
    return correct / n

test_acc_1 = evaluate(model_1, test_image_1, test_label)
print('Testing Accuracy : %.4f' % test_acc_1)

Testing Accuracy : 0.9218

From the above result, we got about 92% accuracy for Softmax Regression on MNIST. In fact, it's not so good. This is because we're using a very simple model.

Multilayer Convolutional Network on MNIST

We're now jumping from a very simple model to something moderately sophisticated: a small Convolutional Neural Network. This will get us to around 99% accuracy, not state of the art, but respectable.

# reshaping the training data to (N, C, H, W) -- PyTorch's Conv2d convention,
# vs. TF/Keras' (N, H, W, C)
train_image_2 = train_image.reshape((60000, 1, 28, 28))
test_image_2 = test_image.reshape((10000, 1, 28, 28))
print(train_image_2.shape)
print(test_image_2.shape)

(60000, 1, 28, 28)
(10000, 1, 28, 28)

Create the convolutional base

As input, a CNN in PyTorch takes tensors of shape (color_channels, image_height, image_width) (channels-first), ignoring the batch size. If you are new to color channels, MNIST has one (because the images are grayscale), whereas a color image has three (R,G,B). In this example, we will configure our CNN to process inputs of shape (1, 28, 28), which is why we reshaped the data above.

# The convolutional base: a stack of Conv2d and MaxPool2d layers (same pattern
# as the Keras Sequential version, just built as an nn.Module)
class ConvBase2(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=0)
        self.pool1 = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=0)
        self.pool2 = nn.MaxPool2d(2, 2)
        self.conv3 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=0)

    def forward(self, x):
        x = self.pool1(F.relu(self.conv1(x)))
        x = self.pool2(F.relu(self.conv2(x)))
        x = F.relu(self.conv3(x))
        return x

Let's display the architecture of our model so far.

conv_base = ConvBase2().to(device)
# there's no built-in `.summary()` in PyTorch; a quick shape probe does the job
with torch.no_grad():
    out = conv_base(torch.zeros(1, 1, 28, 28, device=device))
print(conv_base)
print('output shape:', tuple(out.shape[1:]))

ConvBase2(
  (conv1): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1))
  (pool1): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  (conv2): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1))
  (pool2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  (conv3): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1))
)
output shape: (64, 3, 3)

Above, you can see that the final output of the convolutional base is a 3D tensor of shape (channels, height, width) = (64, 3, 3). The width and height dimensions shrink as we go deeper in the network (28 → 26 → 13 → 11 → 5 → 3), while the number of channels grows. The number of output channels for each Conv2d layer is controlled by its second argument, out_channels (e.g., 32 or 64); the first argument, in_channels, must match the number of channels coming in. Typically, as the width and height shrink, we can afford (computationally) to add more output channels in each Conv2d layer.

Add fully connected layers on top

To complete our model, we will feed the last output tensor from the convolutional base (of shape (64, 3, 3)) into one or more fully connected (nn.Linear) layers to perform classification. Linear layers take vectors as input (which are 1D), while the current output is a 3D tensor. First, we will flatten (or unroll) the 3D output to 1D, then add one or more Linear layers on top. MNIST has 10 output classes, so we use a final Linear layer with 10 outputs. We don't add a softmax at the end: the model outputs raw scores (logits), and nn.CrossEntropyLoss applies the softmax internally during training.

To reduce overfitting, we will apply dropout before the readout layer. The idea behind dropout is to train an ensemble of model instead of a single model. During training, we drop out neurons with probability $p$, i.e., the probability to keep is $1-p$. When a neuron is dropped, its output is set to zero. These dropped neurons do not contribute to the training phase in forward pass and backward pass. For each training phase, we train the network slightly different from the previous one. It's just like we train different networks in each training phrase. However, during testing phase, we don't drop any neuron, and thus, implement dropout is kind of like doing ensemble. Also, randomly drop units in training phase can prevent units from co-adapting too much. Thus, dropout is a powerful regularization techique to deal with overfitting.

class Model2(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv_base = ConvBase2()
        self.flatten = nn.Flatten()
        self.fc1 = nn.Linear(3 * 3 * 64, 64)
        self.dropout = nn.Dropout(0.5)
        self.fc2 = nn.Linear(64, 10)

    def forward(self, x):
        x = self.conv_base(x)
        x = self.flatten(x)
        x = F.relu(self.fc1(x))
        x = self.dropout(x)
        return self.fc2(x)  # logits; softmax folded into CrossEntropyLoss

Here's the complete architecture of our model.

model_2 = Model2().to(device)
print(model_2)
num_params = sum(p.numel() for p in model_2.parameters())
print(f'Total params: {num_params:,}')

Model2(
  (conv_base): ConvBase2(
    (conv1): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1))
    (pool1): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (conv2): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1))
    (pool2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (conv3): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1))
  )
  (flatten): Flatten(start_dim=1, end_dim=-1)
  (fc1): Linear(in_features=576, out_features=64, bias=True)
  (dropout): Dropout(p=0.5, inplace=False)
  (fc2): Linear(in_features=64, out_features=10, bias=True)
)
Total params: 93,322

As you can see, our (64, 3, 3) outputs were flattened into vectors of length 576 (see in_features=576 of fc1) before going through two Linear layers.

Train the model

train_sklearn_style(model_2, train_image_2, train_label, epochs=5)

Epoch 1/5 - loss: 0.4865 - accuracy: 0.8445
Epoch 2/5 - loss: 0.1515 - accuracy: 0.9574
Epoch 3/5 - loss: 0.1023 - accuracy: 0.9718
Epoch 4/5 - loss: 0.0812 - accuracy: 0.9775
Epoch 5/5 - loss: 0.0675 - accuracy: 0.9815

test_acc_2 = evaluate(model_2, test_image_2, test_label)
print('Testing Accuracy : %.4f' % test_acc_2)

Testing Accuracy : 0.9879

As you can see, our simple CNN has achieved a test accuracy of about 99%. Not bad for a few lines of code!

Cifar-10

Actually MNIST is a easy dataset for the beginner. To demonstrate the power of Neural Networks, we need a larger dataset CIFAR-10.

CIFAR-10 consists of 60000 32x32 color images in 10 classes, with 6000 images per class. There are 50000 training images and 10000 test images. Here are the classes in the dataset, as well as 10 random images from each:

Before jumping to a complicated neural network model, we're going to start with KNN and SVM. The motivation here is to compare neural network model with traditional classifiers, and highlight the performance of neural network model.

torchvision.datasets offers convenient facilities that automatically access some well-known datasets. Let's load CIFAR-10:

# Loading Data
cifar_train = torchvision.datasets.CIFAR10(root='./data', train=True, download=True)
cifar_test = torchvision.datasets.CIFAR10(root='./data', train=False, download=True)

X_train = cifar_train.data.astype('float32') / 255.0   # (N, 32, 32, 3), 0..1
y_train = np.array(cifar_train.targets)
X_test = cifar_test.data.astype('float32') / 255.0
y_test = np.array(cifar_test.targets)

# convert class vectors to one-hot (kept only because the original notebook
# used categorical_crossentropy; nn.CrossEntropyLoss wants plain class ids so
# we normally wouldn't one-hot these, but we keep Y_train/Y_test around for
# parity with the original variable names below)
Y_train = np.eye(10, dtype='float32')[y_train]
Y_test = np.eye(10, dtype='float32')[y_test]

print('X_train shape:', X_train.shape)
print('Y_train shape:', Y_train.shape)
print('X_test shape:', X_test.shape)
print('Y_test shape:', Y_test.shape)

X_train shape: (50000, 32, 32, 3)
Y_train shape: (50000, 10)
X_test shape: (10000, 32, 32, 3)
Y_test shape: (10000, 10)

For simplicity, we also convert the images into grayscale. We use the Luma coding that is common in video systems:

# transform a 3-channel image into one channel (unchanged: plain NumPy, no
# TF/PyTorch dependency at all in the original)
def grayscale(data, dtype='float32'):
    # luma coding weighted average in video systems
    r = np.asarray(.3, dtype=dtype)
    g = np.asarray(.59, dtype=dtype)
    b = np.asarray(.11, dtype=dtype)
    rst = r * data[:, :, :, 0] + g * data[:, :, :, 1] + b * data[:, :, :, 2]
    # add channel dimension
    rst = np.expand_dims(rst, axis=3)
    return rst

X_train_gray = grayscale(X_train)
X_test_gray = grayscale(X_test)

# plot a randomly chosen image
img = round(np.random.rand() * X_train.shape[0])
plt.figure(figsize=(4, 2))
plt.subplot(1, 2, 1)
plt.imshow(X_train[img], interpolation='none')
plt.subplot(1, 2, 2)
plt.imshow(
    X_train_gray[img, :, :, 0], cmap=plt.get_cmap('gray'), interpolation='none')
plt.show()

As we can see, the objects in grayscale images can still be recognizable.

first_img = X_train_gray[0,:,:,:]
first_img[0,0]

array([0.24003921], dtype=float32)

Feature Selection

When coming to object detection, HOG (histogram of oriented gradients) is often extracted as a feature for classification. It first calculates the gradients of each image patch using sobel filter, then use the magnitudes and orientations of derived gradients to form a histogram per patch (a vector). After normalizing these histograms, it concatenates them into one HOG feature. For more details, read this tutorial.

Note. one can directly feed the original images for classification; however, it will take lots of time to train and get worse performance.

def getHOGfeat(image,
               stride=8,
               orientations=8,
               pixels_per_cell=(8, 8),
               cells_per_block=(2, 2)):
    cx, cy = pixels_per_cell
    bx, by = cells_per_block
    sx, sy, sz = image.shape
    image = image[:, :, 0]   # strip channel dim so pixel access returns scalars
    n_cellsx = int(np.floor(sx // cx))  # number of cells in x
    n_cellsy = int(np.floor(sy // cy))  # number of cells in y
    n_blocksx = (n_cellsx - bx) + 1
    n_blocksy = (n_cellsy - by) + 1
    gx = np.zeros((sx, sy), dtype=np.double)
    gy = np.zeros((sx, sy), dtype=np.double)
    eps = 1e-5
    grad = np.zeros((sx, sy, 2), dtype=np.double)
    for i in range(1, sx - 1):
        for j in range(1, sy - 1):
            gx[i, j] = image[i, j - 1] - image[i, j + 1]
            gy[i, j] = image[i + 1, j] - image[i - 1, j]
            grad[i, j, 0] = np.arctan(gy[i, j] / (gx[i, j] + eps)) * 180 / math.pi
            if gx[i, j] < 0:
                grad[i, j, 0] += 180
            grad[i, j, 0] = (grad[i, j, 0] + 360) % 360
            grad[i, j, 1] = np.sqrt(gy[i, j] ** 2 + gx[i, j] ** 2)
    normalised_blocks = np.zeros((n_blocksy, n_blocksx, by * bx * orientations))
    for y in range(n_blocksy):
        for x in range(n_blocksx):
            block = grad[y * stride:y * stride + 16, x * stride:x * stride + 16]
            hist_block = np.zeros(32, dtype=np.double)
            eps = 1e-5
            for k in range(by):
                for m in range(bx):
                    cell = block[k * 8:(k + 1) * 8, m * 8:(m + 1) * 8]
                    hist_cell = np.zeros(8, dtype=np.double)
                    for i in range(cy):
                        for j in range(cx):
                            n = int(cell[i, j, 0] / 45)
                            hist_cell[n] += cell[i, j, 1]
                    hist_block[(k * bx + m) * orientations:(k * bx + m + 1) * orientations] = hist_cell[:]
            normalised_blocks[y, x, :] = hist_block / np.sqrt(
                hist_block.sum() ** 2 + eps)
    return normalised_blocks.ravel()

Once we have our getHOGfeat function, we then get the HOG features of all images.

X_train_hog = []
X_test_hog = []

print('This will take some minutes.')

for img in tqdm(X_train_gray):
    img_hog = getHOGfeat(img)
    X_train_hog.append(img_hog)

for img in tqdm(X_test_gray):
    img_hog = getHOGfeat(img)
    X_test_hog.append(img_hog)
    
X_train_hog_array = np.asarray(X_train_hog)
X_test_hog_array = np.asarray(X_test_hog)

This will take some minutes.

100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 50000/50000 [03:33<00:00, 234.13it/s]
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 10000/10000 [00:42<00:00, 234.19it/s]

K Nearest Neighbors (KNN) on CIFAR-10

scikit-learn provides off-the-shelf libraries for classification. For KNN and SVM classifiers, we can just import from scikit-learn to use.

# KNN (scikit-learn -- no TF/PyTorch involved either way)
from sklearn.neighbors import KNeighborsClassifier 
from sklearn.metrics import accuracy_score

# p=2 and metric='minkowski' means the Euclidean Distance
knn = KNeighborsClassifier(n_neighbors=11, p=2, metric='minkowski')

knn.fit(X_train_hog_array, y_train.ravel())
y_pred = knn.predict(X_test_hog_array)
print('[KNN]')
print('Misclassified samples: %d' % (y_test.ravel() != y_pred).sum())
print('Accuracy: %.2f' % accuracy_score(y_test, y_pred))

[KNN]
Misclassified samples: 5334
Accuracy: 0.47

We can observe that the accuracy of KNN on CIFAR-10 is embarrassingly bad.

Support Vector Machine (SVM) on CIFAR-10

# SVM
from sklearn.svm import SVC 

print('This will take some minutes.')
start_time = time.time()

# C is the hyperparameter for the error penalty term
# gamma is the hyperparameter for the rbf kernel
svm_linear = SVC(kernel='linear', random_state=0, gamma=0.2, C=10.0)

svm_linear.fit(X_train_hog_array, y_train.ravel())
y_pred = svm_linear.predict(X_test_hog_array)
print('[Linear SVC]')
print('Misclassified samples: %d' % (y_test.ravel() != y_pred).sum())
print('Accuracy: %.2f' % accuracy_score(y_test.ravel(), y_pred))

print('{:.2f} sec.'.format(time.time()-start_time))

This will take some minutes.
[Linear SVC]
Misclassified samples: 4940
Accuracy: 0.51
537.22 sec.

By above, SVM is slightly better than KNN, but still poor. Next, we'll design a CNN model using PyTorch.

CNN on CIFAR-10

class CifarCNN(nn.Module):
    """Same architecture as the Keras `model_3`: 2x (Conv-Pool-BN) + MLP head."""
    def __init__(self, in_hw=32):
        super().__init__()
        # padding='same' + stride 1 for the 5x5 convs -> padding=2
        self.conv1 = nn.Conv2d(3, 64, kernel_size=5, padding=2)
        self.pool1 = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)  # padding='same'
        self.bn1 = nn.BatchNorm2d(64)
        self.conv2 = nn.Conv2d(64, 64, kernel_size=5, padding=2)
        self.pool2 = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        self.bn2 = nn.BatchNorm2d(64)

        flat_hw = math.ceil(math.ceil(in_hw / 2) / 2)
        self.flatten = nn.Flatten()
        self.fc1 = nn.Linear(64 * flat_hw * flat_hw, 384)
        self.dropout = nn.Dropout(0.5)
        self.fc2 = nn.Linear(384, 192)
        self.fc3 = nn.Linear(192, 10)

    def forward(self, x):
        x = self.bn1(self.pool1(F.relu(self.conv1(x))))
        x = self.bn2(self.pool2(F.relu(self.conv2(x))))
        x = self.flatten(x)
        x = F.relu(self.fc1(x))
        x = self.dropout(x)
        x = F.relu(self.fc2(x))
        return self.fc3(x)

model_3 = CifarCNN().to(device)
print(model_3)

CifarCNN(
  (conv1): Conv2d(3, 64, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
  (pool1): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
  (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (conv2): Conv2d(64, 64, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
  (pool2): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
  (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (flatten): Flatten(start_dim=1, end_dim=-1)
  (fc1): Linear(in_features=4096, out_features=384, bias=True)
  (dropout): Dropout(p=0.5, inplace=False)
  (fc2): Linear(in_features=384, out_features=192, bias=True)
  (fc3): Linear(in_features=192, out_features=10, bias=True)
)

# NHWC -> NCHW for PyTorch's Conv2d
X_train_chw = np.transpose(X_train, (0, 3, 1, 2))
X_test_chw = np.transpose(X_test, (0, 3, 1, 2))

train_sklearn_style(model_3, X_train_chw, y_train, epochs=5)
test_acc_3 = evaluate(model_3, X_test_chw, y_test)
print('Testing Accuracy : %.4f' % test_acc_3)

Epoch 1/5 - loss: 0.5208 - accuracy: 0.8167
Epoch 2/5 - loss: 0.4658 - accuracy: 0.8369
Epoch 3/5 - loss: 0.4139 - accuracy: 0.8558
Epoch 4/5 - loss: 0.3671 - accuracy: 0.8719
Epoch 5/5 - loss: 0.3305 - accuracy: 0.8835
Testing Accuracy : 0.7704

Although Cifar10 is larger than Mnist, it's not large enough for the dataset you will meet in the following lessons. For large datasets, we can't feed all training data to the model due to the limit of GPU memory size. Even if we can feed all data into the model, we still want the process of loading data is efficient. Input pipeline is the common way to solve these.

Input Pipeline

Structure of a data loading pipeline

A typical training data pipeline can be framed as an ETL process:

1. Extract: Read data from memory (NumPy) or persistent storage -- either local (HDD/SSD) or remote.
2. Transform: parse and preprocess the data -- shuffling, batching, domain-specific transformations such as image decoding/augmentation.
3. Load: move the transformed data onto the accelerator device(s) that execute the model.

This pattern keeps the CPU busy preparing data while reserving the accelerator for the heavy lifting of training your model.

In PyTorch, this job is split across two plain-Python objects:

• torch.utils.data.Dataset: a class you write, exposing __len__ and __getitem__(idx) -- this is the "Extract + per-sample Transform" step. Preprocessing is done inside __getitem__, often by calling a torchvision.transforms pipeline.
• torch.utils.data.DataLoader: wraps a Dataset and handles shuffling, batching, parallel loading (num_workers), and prefetching -- this is the "Load" step. All of these are configured as constructor arguments.

Construct your Dataset

Suppose we have simple data sources:

# number of samples
n_samples = 200

# an array with shape (n_samples, 5)
raw_data_a = np.random.rand(n_samples, 5).astype('float32')
# a list with length of n_samples from 0 to n_samples-1
raw_data_b = np.arange(n_samples)
print(raw_data_a.shape, raw_data_b.shape)

(200, 5) (200,)

We wrap these two arrays in a Dataset, which pairs each row of raw_data_a with the corresponding element of raw_data_b and returns one sample per index. We also add an optional transform hook, which we'll use for preprocessing further down:

class ArrayDataset(Dataset):
    # this tells the dataset that each row of `a` corresponds to each element of `b`
    def __init__(self, a, b, transform=None):
        self.a, self.b = a, b
        self.transform = transform

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

    def __getitem__(self, idx):
        a, b = self.a[idx], self.b[idx]
        if self.transform is not None:
            a, b = self.transform(a, b)
        return a, b

raw_dataset = ArrayDataset(raw_data_a, raw_data_b)

Consume elements

A map-style Dataset (one that implements __getitem__/__len__) is a Python iterable. This makes it possible to consume its elements using a for loop:

# Here, we print the first 8 samples.
for i, elem in enumerate(raw_dataset):
    print("Sample ", i, ", b is ", elem[1])
    if i == 7:
        break

Sample  0 , b is  0
Sample  1 , b is  1
Sample  2 , b is  2
Sample  3 , b is  3
Sample  4 , b is  4
Sample  5 , b is  5
Sample  6 , b is  6
Sample  7 , b is  7

Or by explicitly creating a Python iterator using iter and consuming its elements using next:

# Here, we print the first 8 samples.
it = iter(raw_dataset)
for i in range(8):
    print("Sample ", i, ", b is ", next(it)[1])

Sample  0 , b is  0
Sample  1 , b is  1
Sample  2 , b is  2
Sample  3 , b is  3
Sample  4 , b is  4
Sample  5 , b is  5
Sample  6 , b is  6
Sample  7 , b is  7

Apply transformations

map

In PyTorch, element-wise preprocessing lives inside __getitem__ -- here, through the transform hook we already built in. Let's set it to a function that transforms each row of a into its sum and mean:

def preprocess_function(one_row_a, one_b):
    """
        Input: one slice of the dataset
        Output: modified slice
    """
    # Do some data preprocessing, you can also input filenames and load data in here
    # Here, we transform each row of raw_data_a to its sum and mean
    one_row_a = [np.sum(one_row_a), np.mean(one_row_a)]
    return one_row_a, one_b

raw_dataset.transform = preprocess_function

it = iter(raw_dataset)
for i in range(8):
    print("Sample ", i, ", b is ", next(it))

Sample  0 , b is  ([np.float32(2.893259), np.float32(0.5786518)], np.int64(0))
Sample  1 , b is  ([np.float32(2.6859148), np.float32(0.5371829)], np.int64(1))
Sample  2 , b is  ([np.float32(2.2667904), np.float32(0.45335808)], np.int64(2))
Sample  3 , b is  ([np.float32(2.6292167), np.float32(0.5258433)], np.int64(3))
Sample  4 , b is  ([np.float32(2.3166103), np.float32(0.46332207)], np.int64(4))
Sample  5 , b is  ([np.float32(3.0557249), np.float32(0.61114496)], np.int64(5))
Sample  6 , b is  ([np.float32(2.2304804), np.float32(0.4460961)], np.int64(6))
Sample  7 , b is  ([np.float32(3.4401903), np.float32(0.68803805)], np.int64(7))

shuffle, batch, repeat, prefetch with DataLoader

Shuffling, batching, iterating over multiple epochs, and prefetching are all handled by the DataLoader that wraps the Dataset:

• shuffle: DataLoader(shuffle=True) uses a RandomSampler that reshuffles the indices of the entire dataset at every epoch. This way, you see your data in a different order in each epoch, which prevents your model from overfitting to the order of your training data (it assumes a map-style Dataset with a known __len__, which is the common case).

loader = DataLoader(raw_dataset, batch_size=1, shuffle=True)

idxs = []
for i, elem in enumerate(loader):
    print("Sample ", i, ", b is ", elem[1])
    idxs.append(elem[1].item())
    if i == 7:
        break

print("\nThe order of the first 8 shuffled from [0, 1, 2, 3, 4, 5, 6, 7] to ", idxs)

Sample  0 , b is  tensor([124])
Sample  1 , b is  tensor([77])
Sample  2 , b is  tensor([65])
Sample  3 , b is  tensor([164])
Sample  4 , b is  tensor([75])
Sample  5 , b is  tensor([39])
Sample  6 , b is  tensor([111])
Sample  7 , b is  tensor([160])

The order of the first 8 shuffled from [0, 1, 2, 3, 4, 5, 6, 7] to  [124, 77, 65, 164, 75, 39, 111, 160]

• batch: Our dataset returns one example at a time. However, in practice we usually want to read one batch at a time, so we set batch_size= to stack that many samples together. drop_last=True drops the final incomplete batch (the default is drop_last=False).

  Note: DataLoader always shuffles sample indices first (via its internal RandomSampler) and then groups them into batches (via BatchSampler), so each batch contains randomly chosen samples rather than a fixed group of neighboring samples.

loader = DataLoader(raw_dataset, batch_size=2, shuffle=True, drop_last=False)

idxs = []
for i, elem in enumerate(loader):
    print("Batch ", i, ", b are ", elem[1])
    print("")
    idxs.append(elem[1].tolist())
    if i == 7:
        break

print("After `DataLoader(batch_size=2)`,\nBatch 0 is combined by %dth and %dth,\nBatch 1 is combined by %dth and %dth,\nBatch 2 is combined by %dth and %dth, etc."
      % (idxs[0][0], idxs[0][1], idxs[1][0], idxs[1][1], idxs[2][0], idxs[2][1]))

Batch  0 , b are  tensor([111,  54])

Batch  1 , b are  tensor([17, 96])

Batch  2 , b are  tensor([63, 74])

Batch  3 , b are  tensor([ 2, 53])

Batch  4 , b are  tensor([ 26, 101])

Batch  5 , b are  tensor([31, 94])

Batch  6 , b are  tensor([66, 95])

Batch  7 , b are  tensor([147, 133])

After `DataLoader(batch_size=2)`,
Batch 0 is combined by 111th and 54th,
Batch 1 is combined by 17th and 96th,
Batch 2 is combined by 63th and 74th, etc.

• repeat: There is no .repeat() in PyTorch. A "new epoch" is just another pass over the loader, i.e., calling iter(loader) again (which the outer for epoch in range(epochs): for batch in loader: loop does for you). Every epoch is a clean, complete pass with a fresh (re-shuffled, if shuffle=True) ordering, so batches never straddle epoch boundaries. This also makes it simple to perform a custom computation (e.g. to collect statistics) at the end of each epoch:

epochs = 3

for epoch in range(epochs):
    size = 0
    n_batch = 0
    for batch in loader:
        size += len(batch[1])
        n_batch += 1
    print("End of epoch %d: Total %d batches in this epoch with %d pieces of data" % (epoch, n_batch, size))

End of epoch 0: Total 100 batches in this epoch with 200 pieces of data
End of epoch 1: Total 100 batches in this epoch with 200 pieces of data
End of epoch 2: Total 100 batches in this epoch with 200 pieces of data

• prefetch: Prefetching overlaps producing the next batch with consuming the current one. The relevant DataLoader arguments are:

  • num_workers > 0: spawn background worker processes that load/transform samples in parallel, so upcoming batches are prepared while the model trains on the current one.
  • prefetch_factor: how many batches each worker prepares ahead of time (default 2 per worker when num_workers > 0).
  • pin_memory=True: allocate batches in page-locked host memory so the host-to-GPU copy (.to(device)) can happen asynchronously/faster.

  We'll actually measure the effect of these in the benchmark below.

loader = DataLoader(raw_dataset, batch_size=2, shuffle=True,
                     num_workers=2, prefetch_factor=2, pin_memory=torch.cuda.is_available())

Now, let's start designing our CNN model!

CNN Model for CIFAR 10

Loading Data Manually

To know how it works under the hood, let's load CIFAR-10 by our own (not using torchvision.datasets.CIFAR10). According to the description, the dataset file is divided into five training batches and one test batch, each with 10000 images. The test batch contains exactly 1000 randomly-selected images from each class.

# the url to download CIFAR-10 dataset (binary version)
# see format and details here: http://www.cs.toronto.edu/~kriz/cifar.html
DATA_URL = 'https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz'
# the image size we want to keep
IMAGE_SIZE_CROPPED = 24
IMAGE_HEIGHT = 32
IMAGE_WIDTH = 32
IMAGE_DEPTH = 3

import urllib.request
import tarfile

# download data
if not os.path.exists("cifar-10-batches-py/"):
    archive_path, _ = urllib.request.urlretrieve(DATA_URL, 'cifar-10-python.tar.gz')
    with tarfile.open(archive_path) as f:
        f.extractall('.')

if not os.path.exists("cifar-10-batches-py/"):
    with tarfile.open('cifar-10-python.tar.gz') as f:
        f.extractall('.')

DEST_DIRECTORY = 'cifar-10-batches-py'
filenames_train = [os.path.join(DEST_DIRECTORY,'data_batch_%d' % i) for i in range(1,6)]
filenames_test = [os.path.join(DEST_DIRECTORY,'test_batch')]

# save (img_path, label) pairs
with open('cifar10_train.csv', 'w', newline='') as csvfile:
    writer = csv.writer(csvfile)
    writer.writerow(['filenames'])
    writer.writerows(np.array(filenames_train).reshape(-1,1))
        
with open('cifar10_test.csv', 'w', newline='') as csvfile:
    writer = csv.writer(csvfile)
    writer.writerow(['filenames'])
    writer.writerows(np.array(filenames_test).reshape(-1,1))

Now we implement the per-sample preprocessing (reshape, normalize, random crop, random flip, brightness/contrast jitter, per-image standardization) as a torchvision.transforms pipeline, and a Dataset that reads the raw pickled CIFAR batches. Note that the random augmentations are only applied to the training set; the test set only gets a center crop and standardization:

# read data
def read_file(file):
    with open(file, 'rb') as fo:
        raw_data = pickle.load(fo, encoding='bytes')
    return raw_data[b'data'], raw_data[b'labels']

class CifarPickleDataset(Dataset):
    def __init__(self, filenames, train):
        images, labels = [], []
        for filename in filenames:
            image, label = read_file(filename)
            images.append(image)
            labels.extend(label)
        self.images = np.concatenate(images).reshape(-1, IMAGE_DEPTH, IMAGE_HEIGHT, IMAGE_WIDTH)
        self.labels = np.array(labels, dtype=np.int64)

        # equivalent of map_fun / map_fun_test
        aug = [transforms.RandomCrop(IMAGE_SIZE_CROPPED)] if train else \
              [transforms.CenterCrop(IMAGE_SIZE_CROPPED)]
        if train:
            aug += [
                transforms.RandomHorizontalFlip(),
                transforms.ColorJitter(brightness=63/255, contrast=(0.2, 1.8)),
            ]
        self.transform = transforms.Compose(aug + [
            # per-image standardization: TF's per_image_standardization()
            # subtracts the per-image mean and divides by (approx.) the per-image std
            transforms.Lambda(lambda t: (t - t.mean()) / t.std().clamp_min(1.0 / math.sqrt(t.numel()))),
        ])

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

    def __getitem__(self, idx):
        image = torch.from_numpy(self.images[idx]).float() / 255.0  # already CHW
        image = self.transform(image)
        return image, self.labels[idx]

# Construct training/testing Dataset + DataLoader
train_dataset = CifarPickleDataset(filenames_train, train=True)
test_dataset = CifarPickleDataset(filenames_test, train=False)

train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=64)

model_cifar = CifarCNN(in_hw=IMAGE_SIZE_CROPPED).to(device)
print(model_cifar)

CifarCNN(
  (conv1): Conv2d(3, 64, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
  (pool1): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
  (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (conv2): Conv2d(64, 64, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
  (pool2): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
  (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (flatten): Flatten(start_dim=1, end_dim=-1)
  (fc1): Linear(in_features=2304, out_features=384, bias=True)
  (dropout): Dropout(p=0.5, inplace=False)
  (fc2): Linear(in_features=384, out_features=192, bias=True)
  (fc3): Linear(in_features=192, out_features=10, bias=True)
)

def train_loader_style(model, train_loader, test_loader, epochs, lr=1e-3):
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    criterion = nn.CrossEntropyLoss()

    for epoch in range(epochs):
        model.train()
        total_loss, correct, n = 0.0, 0, 0
        for xb, yb in train_loader:
            xb, yb = xb.to(device), yb.to(device)
            optimizer.zero_grad()
            logits = model(xb)
            loss = criterion(logits, yb)
            loss.backward()
            optimizer.step()
            total_loss += loss.item() * xb.size(0)
            correct += (logits.argmax(1) == yb).sum().item()
            n += xb.size(0)

        val_acc = evaluate_loader(model, test_loader)
        print(f'Epoch {epoch+1}/{epochs} - loss: {total_loss/n:.4f} '
              f'- accuracy: {correct/n:.4f} - val_accuracy: {val_acc:.4f}')

@torch.no_grad()
def evaluate_loader(model, loader):
    model.eval()
    correct, n = 0, 0
    for xb, yb in loader:
        xb, yb = xb.to(device), yb.to(device)
        correct += (model(xb).argmax(1) == yb).sum().item()
        n += xb.size(0)
    return correct / n

train_loader_style(model_cifar, train_loader, test_loader, epochs=5)
test_acc = evaluate_loader(model_cifar, test_loader)
print('test accuracy:', test_acc)

Epoch 1/5 - loss: 1.4615 - accuracy: 0.4707 - val_accuracy: 0.6196
Epoch 2/5 - loss: 1.1798 - accuracy: 0.5830 - val_accuracy: 0.6669
Epoch 3/5 - loss: 1.0621 - accuracy: 0.6259 - val_accuracy: 0.6768
Epoch 4/5 - loss: 0.9981 - accuracy: 0.6522 - val_accuracy: 0.7151
Epoch 5/5 - loss: 0.9539 - accuracy: 0.6688 - val_accuracy: 0.7261
test accuracy: 0.7261

Optimization for input pipeline

We all know that GPUs can radically reduce the time required to execute a single training step; however, all other affairs (data loading, transformations, memory copy from CPU to GPU) are done by CPU, which sometimes becomes the bottleneck instead. We are going to accelerate the data loading pipeline for better training performance using PyTorch's own knobs, following the DataLoader performance guide.

The code below re-runs roughly the same CIFAR-10 training as above, but we add an artificial I/O latency (time.sleep) to every sample read -- simulating a slow disk/network dataset -- purely so the Naive vs. Optimized difference is clearly visible on a small, already-in-memory dataset like CIFAR-10.

A timing-instrumented Dataset

TimeMeasuredDataset (in lab11_1_lib_torch.py) is a Dataset that records (Open, Read, Map) wall-clock timings per sample. The training loop below appends matching Train/Test timing entries, and we visualize everything with draw_timeline().

def open_cifar(filenames):
    images, labels = [], []
    for filename in filenames:
        image, label = read_file(filename)
        images.append(image.reshape(-1, IMAGE_DEPTH, IMAGE_HEIGHT, IMAGE_WIDTH))
        labels.extend(label)
    images = np.concatenate(images)
    return list(images), np.array(labels, dtype=np.int64)

def decode_cifar(raw_image):
    return torch.from_numpy(raw_image).float() / 255.0

def build_transform(train):
    aug = [transforms.RandomCrop(IMAGE_SIZE_CROPPED)] if train else \
          [transforms.CenterCrop(IMAGE_SIZE_CROPPED)]
    if train:
        aug += [transforms.RandomHorizontalFlip(),
                transforms.ColorJitter(brightness=63/255, contrast=(0.2, 1.8))]
    aug += [transforms.Lambda(lambda t: (t - t.mean()) / t.std().clamp_min(1.0 / math.sqrt(t.numel())))]
    return transforms.Compose(aug)

from lab11_1_lib_torch import TimeMeasuredDataset, timing_collate

IO_LATENCY = 0.003  # seconds, simulated per-sample disk read

def make_timed_datasets():
    train_ds = TimeMeasuredDataset(open_cifar, decode_cifar, filenames_train,
                                    transform=build_transform(train=True),
                                    simulate_io_latency=IO_LATENCY)
    test_ds = TimeMeasuredDataset(open_cifar, decode_cifar, filenames_test,
                                   transform=build_transform(train=False),
                                   simulate_io_latency=IO_LATENCY)
    return train_ds, test_ds

def timelined_benchmark(model, train_loader, test_loader, epochs, lr=1e-3):
    steps_acc, times_acc, values_acc = [], [], []
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    criterion = nn.CrossEntropyLoss()

    start_time = time.perf_counter()
    print("start time: ", start_time)
    for epoch in range(epochs):
        model.train()
        print("training:")
        for steps, times, values, image, label in tqdm(train_loader):
            for s in steps: steps_acc.extend(s)
            for t in times: times_acc.extend(t)
            for v in values: values_acc.extend(v)

            image, label = image.to(device, non_blocking=True), label.to(device, non_blocking=True)
            train_enter = time.perf_counter()
            optimizer.zero_grad()
            loss = criterion(model(image), label)
            loss.backward()
            optimizer.step()
            train_elapsed = time.perf_counter() - train_enter

            steps_acc.append('Train')
            times_acc.append((train_enter, train_elapsed))
            values_acc.append((-1, epoch, -1))

        model.eval()
        print("testing:")
        with torch.no_grad():
            for steps, times, values, image, label in tqdm(test_loader):
                for s in steps: steps_acc.extend(s)
                for t in times: times_acc.extend(t)
                for v in values: values_acc.extend(v)

                image, label = image.to(device), label.to(device)
                test_enter = time.perf_counter()
                model(image)
                test_elapsed = time.perf_counter() - test_enter
                steps_acc.append('Test')
                times_acc.append((test_enter, test_elapsed))
                values_acc.append((-1, epoch, -1))

    print("Execution time:", time.perf_counter() - start_time)
    return {
        'steps': np.array(steps_acc),
        'times': np.array(times_acc, dtype=np.float64),
        'values': np.array(values_acc, dtype=np.int64),
    }

Naive: num_workers=0, no pin_memory, no caching

train_ds, test_ds = make_timed_datasets()
train_loader_naive = DataLoader(train_ds, batch_size=64, shuffle=True,
                                 num_workers=0, collate_fn=timing_collate)
test_loader_naive = DataLoader(test_ds, batch_size=64,
                                num_workers=0, collate_fn=timing_collate)

model_cifar_op = CifarCNN(in_hw=IMAGE_SIZE_CROPPED).to(device)
torch.save(model_cifar_op.state_dict(), 'model_cifar_op_init.pt')  # save initial weights

timeline_naive = timelined_benchmark(model_cifar_op, train_loader_naive, test_loader_naive, epochs=2)

start time:  354196.876879551
training:

100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 782/782 [03:42<00:00,  3.52it/s]

testing:

100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 157/157 [00:36<00:00,  4.33it/s]

training:

100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 782/782 [03:41<00:00,  3.52it/s]

testing:

100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 157/157 [00:36<00:00,  4.33it/s]

Execution time: 516.5708793090307

draw_timeline(timeline=timeline_naive, title="Naive", min_width=1)

Optimized: num_workers, pin_memory, prefetch_factor + caching

The optimizations we apply:

1. Parallel loading -> num_workers > 0: multiple worker processes read/decode samples concurrently.
2. Prefetching -> prefetch_factor + pin_memory=True: batches are prepared ahead of time in page-locked memory, so the host->GPU copy overlaps with compute.
3. Caching -> a small in-Dataset cache of the decoded image (before the random augmentation transform). We deliberately cache before augmentation: caching the augmented result would silently freeze the random augmentation to whatever was computed on the first epoch -- a common gotcha. Caching pre-augmentation avoids that while still skipping the (simulated) slow read on repeat epochs.

class CachedTimeMeasuredDataset(TimeMeasuredDataset):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._decoded_cache = {}

    def __getitem__(self, idx):
        epoch_idx, sample_idx = divmod(self._call_count, len(self))
        self._call_count += 1

        if idx in self._decoded_cache:
            read_enter = time.perf_counter()
            img = self._decoded_cache[idx]
            read_elapsed = time.perf_counter() - read_enter  # ~0, cache hit
        else:
            read_enter = time.perf_counter()
            if self.simulate_io_latency:
                time.sleep(self.simulate_io_latency)
            img = self.read_file(self.paths[idx])
            read_elapsed = time.perf_counter() - read_enter
            self._decoded_cache[idx] = img

        map_enter = time.perf_counter()
        img = self.transform(img) if self.transform is not None else img
        map_elapsed = time.perf_counter() - map_enter

        if not self._opened_once:
            open_entry = self._open_timing
            self._opened_once = True
        else:
            open_entry = (-1.0, -1.0)

        steps = ['Open', 'Read', 'Map']
        times = [open_entry, (read_enter, read_elapsed), (map_enter, map_elapsed)]
        values = [(self.instance_idx, epoch_idx, -1),
                  (self.instance_idx, epoch_idx, sample_idx),
                  (self.instance_idx, epoch_idx, sample_idx)]
        return steps, times, values, img, self.labels[idx]

train_ds_opt = CachedTimeMeasuredDataset(open_cifar, decode_cifar, filenames_train,
                                          transform=build_transform(train=True),
                                          simulate_io_latency=IO_LATENCY)
test_ds_opt = CachedTimeMeasuredDataset(open_cifar, decode_cifar, filenames_test,
                                         transform=build_transform(train=False),
                                         simulate_io_latency=IO_LATENCY)

train_loader_opt = DataLoader(train_ds_opt, batch_size=64, shuffle=True,
                               num_workers=4, pin_memory=torch.cuda.is_available(),
                               prefetch_factor=4, collate_fn=timing_collate)
test_loader_opt = DataLoader(test_ds_opt, batch_size=64,
                              num_workers=2, pin_memory=torch.cuda.is_available(),
                              collate_fn=timing_collate)

# load the same initialization of weights and re-train with the optimized loader
model_cifar_op.load_state_dict(torch.load('model_cifar_op_init.pt'))
timeline_optimized = timelined_benchmark(model_cifar_op, train_loader_opt, test_loader_opt, epochs=2)

start time:  354776.099262812
training:

100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 782/782 [00:54<00:00, 14.24it/s]

testing:

100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 157/157 [00:18<00:00,  8.63it/s]

training:

100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 782/782 [00:55<00:00, 14.21it/s]

testing:

100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 157/157 [00:18<00:00,  8.63it/s]

Execution time: 146.3372294830042

draw_timeline(timeline_optimized, "Optimized", min_width=1)

Comparing the two timelines, the total execution time drops noticeably (from about 517 to 146 seconds in this run), mainly because the worker processes overlap Read/Map with Train.

Note that caching does not help in this run: epoch 1 and epoch 2 take about the same time. With num_workers > 0, each worker process holds its own copy of the Dataset (and therefore its own cache), and by default the workers are shut down and restarted at every epoch, so the cache is thrown away. Passing persistent_workers=True to the DataLoader keeps the workers -- and their caches -- alive across epochs.

The exact numbers are highly device-dependent (CPU core count, disk speed, whether you're on GPU/MPS/CPU) -- re-running this on your own machine may look different, especially once training itself (not I/O) is the bottleneck.

num_workers is almost always worth setting > 0 on a real dataset (CPU core count is a reasonable starting point); remember that pin_memory=True only helps when you're actually training on a CUDA GPU.

Assignment

In this assignment, you have to implement the data loading pipeline of the CNN model with the Oregon Wildlife dataset, and (optionally) try packing the dataset into a single-file archive (see the note at the end).

We provide you with the complete code for the image classification task of the CNN model, but remove the part of the input pipeline. What you need to do is completing this part and training the model for at least 3 epochs.

Description of Dataset:

1. The raw data is from kaggle, which consists of 20 classes of wildlife images.
2. We have filtered the raw data. You need to download the filtered images from here and use them to complete the image classification task.
3. In the dataset we prepared for you, there are nearly 7,200 images, which contain 10 kinds of wildlife.

The sample image is shown below:

red_fox

Requirement:

• Try some of the input transformations mentioned above (e.g. shuffle via DataLoader(shuffle=True), batch_size, transforms.RandomCrop/RandomHorizontalFlip/...) but without the optimization terms (num_workers, pin_memory, prefetch_factor, caching).

  • Note that it's important to look at a few samples and plot the image to make sure the transform does what you intended.

• Retrain your model with the optimized DataLoader settings, comparing the performance and time consumption.

• Train both models above for at least 3 epochs.

• Briefly summarize what you did and explain the performance results (accuracy and time consumption).

  • It's fine if you get worse performance after applying the input transformation, or a longer time after applying the pipeline optimization. Just try to analyze the result and propose your assumption.

• submit your homework to eeclass before 2026/10/28 23:59

• Name your ipynb with the format Lab11_{student_id}.ipynb.

import os
import warnings
warnings.filterwarnings("ignore")

import torch
import torch.nn.functional as F
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import torchvision
from torchvision import transforms
from torchvision.models import vgg16, VGG16_Weights
from sklearn.utils import shuffle
from sklearn.model_selection import train_test_split
import IPython.display as display
import matplotlib.pyplot as plt
import pathlib
import random
import numpy as np
import time

import csv
import pandas as pd
import math
from tqdm import tqdm

from lab11_1_lib_torch import TimeMeasuredDataset, timing_collate, draw_timeline

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print('Using device:', device)

Using device: cuda

# You need to download the prepared data and unzip the file in current path('./')
data_root = pathlib.Path('./oregon_wildlife')

# print the subfolders.
print('classes:')
for item in data_root.iterdir():
    print(item)
    
all_image_paths = list(data_root.glob('*/*'))
all_image_paths = [str(path) for path in all_image_paths]
all_image_paths = shuffle(all_image_paths, random_state=1)
all_image_paths = [path for path in all_image_paths if path[-3:] not in ('gif','bmp')]
image_count = len(all_image_paths)
print('\ntotal img num:', image_count)

classes:
oregon_wildlife/virginia_opossum
oregon_wildlife/raccoon
oregon_wildlife/.DS_Store
oregon_wildlife/red_fox
oregon_wildlife/bald_eagle
oregon_wildlife/black_bear
oregon_wildlife/raven
oregon_wildlife/nutria
oregon_wildlife/sea_lions
oregon_wildlife/deer
oregon_wildlife/cougar

total img num: 7168

# random showing 3 iamges for you
for n in range(3):
    image_path = random.choice(all_image_paths)
    display.display(display.Image(image_path, width=200, height=200))
    print(image_path.split('/')[-2])

cougar

deer

cougar

# get the label
label_names = sorted(item.name for item in data_root.glob('*/') if item.is_dir())
# total label
n_classes = len(label_names)
print(label_names)

['bald_eagle', 'black_bear', 'cougar', 'deer', 'nutria', 'raccoon', 'raven', 'red_fox', 'sea_lions', 'virginia_opossum']

# get the mapping dict
label_to_index = dict((name, index) for index,name in enumerate(label_names))
index_to_label = dict((index, name) for index,name in enumerate(label_names))
print(label_to_index)

{'bald_eagle': 0, 'black_bear': 1, 'cougar': 2, 'deer': 3, 'nutria': 4, 'raccoon': 5, 'raven': 6, 'red_fox': 7, 'sea_lions': 8, 'virginia_opossum': 9}

# get the label data
all_image_label = [label_to_index[pathlib.Path(path).parent.name] for path in all_image_paths]
print("First 10 label indices: ", all_image_label[:10])

First 10 label indices:  [3, 6, 7, 6, 9, 7, 5, 9, 5, 9]

# Create training and testing sets using an 80-20 split
img_path_train, img_path_test, label_train, label_test = train_test_split(all_image_paths,
                                all_image_label,test_size=0.2,random_state=0)
print('training data: %d'%(len(img_path_train)))
print('testing data: %d'%(len(img_path_test)))

training data: 5734
testing data: 1434

# save (img_path, label) pairs
with open('train.csv', 'w', newline='') as csvfile:
    writer = csv.writer(csvfile)
    writer.writerow(['img_path', 'label'])
    for img_path, label in zip(img_path_train, label_train):
        writer.writerow([img_path, label])
        
with open('test.csv', 'w', newline='') as csvfile:
    writer = csv.writer(csvfile)
    writer.writerow(['img_path', 'label'])
    for img_path, label in zip(img_path_test, label_test):
        writer.writerow([img_path, label])

# Feel free to change IMAGE_SIZE_CROPPED if using RandomCrop in your data
# augmentation process, but make sure the input is resized back to
# (224, 224, 3) before feeding into VGG16
IMAGE_SIZE_CROPPED = 224
IMAGE_HEIGHT = 300
IMAGE_WIDTH = 300
IMAGE_DEPTH = 3

Reuse TimeMeasuredDataset from lab11_1_lib_torch.py (same class we used for the CIFAR-10 benchmark above) -- here read_image_file does real per-file I/O (PIL.Image.open + resize), so you'll see genuine disk-read cost in the Read row, not a simulated one.

def open_csv(filename):
    rows = pd.read_csv(filename)
    return rows['img_path'].tolist(), rows['label'].tolist()

def read_image_file(image_path):
    from PIL import Image
    img = Image.open(image_path).convert('RGB').resize((IMAGE_WIDTH, IMAGE_HEIGHT))
    img = torch.from_numpy(np.array(img)).permute(2, 0, 1).float() / 255.0
    return img

def dataset_generator_fun_train():
    return TimeMeasuredDataset(open_csv, read_image_file, 'train.csv',
                                transform=transforms.CenterCrop(IMAGE_SIZE_CROPPED))

def dataset_generator_fun_test():
    return TimeMeasuredDataset(open_csv, read_image_file, 'test.csv',
                                transform=transforms.CenterCrop(IMAGE_SIZE_CROPPED))

# feel free to modify these two settings.
BUFFER_SIZE = 10000
BATCH_SIZE = 32

dataset_train = dataset_generator_fun_train()
dataset_test = dataset_generator_fun_test()

train_loader = DataLoader(dataset_train, batch_size=BATCH_SIZE, shuffle=True,
                           num_workers=0, collate_fn=timing_collate)
test_loader = DataLoader(dataset_test, batch_size=BATCH_SIZE,
                          num_workers=0, collate_fn=timing_collate)

steps, times, values, img, label = dataset_train[0]
print(steps, times, values[1])
print(img.shape)
plt.imshow(img.permute(1, 2, 0))
plt.axis('off')
plt.show()
print(index_to_label[label])

['Open', 'Read', 'Map'] [(1370164.596371635, 0.016556889982894063), (1370165.07124026, 0.031152778072282672), (1370165.10239452, 0.0003703299444168806)] (0, 0, 0)
torch.Size([3, 224, 224])

nutria

base_model = vgg16(weights=VGG16_Weights.IMAGENET1K_V1)
base_model = base_model.features  # convolutional feature extractor only, no top
for p in base_model.parameters():
    p.requires_grad_(False)

class WildModel(nn.Module):
    def __init__(self, n_classes):
        super().__init__()
        self.base_model = base_model
        self.pool = nn.AdaptiveAvgPool2d((7, 7))  # VGG16's expected feature-map size
        self.flatten = nn.Flatten()
        self.fc1 = nn.Linear(512 * 7 * 7, 4096)
        self.drop1 = nn.Dropout(0.5)
        self.fc2 = nn.Linear(4096, 1024)
        self.drop2 = nn.Dropout(0.5)
        self.fc3 = nn.Linear(1024, n_classes)

    def forward(self, x):
        with torch.no_grad():
            x = self.base_model(x)
        x = self.pool(x)
        x = self.flatten(x)
        x = F.relu(self.fc1(x))
        x = self.drop1(x)
        x = F.relu(self.fc2(x))
        x = self.drop2(x)
        return self.fc3(x)

wild_model = WildModel(n_classes).to(device)

print(wild_model)

WildModel(
  (base_model): Sequential(
    (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): ReLU(inplace=True)
    (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (3): ReLU(inplace=True)
    (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (6): ReLU(inplace=True)
    (7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (8): ReLU(inplace=True)
    (9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (11): ReLU(inplace=True)
    (12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (13): ReLU(inplace=True)
    (14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (15): ReLU(inplace=True)
    (16): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (17): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (18): ReLU(inplace=True)
    (19): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (20): ReLU(inplace=True)
    (21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (22): ReLU(inplace=True)
    (23): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (24): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (25): ReLU(inplace=True)
    (26): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (27): ReLU(inplace=True)
    (28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (29): ReLU(inplace=True)
    (30): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (pool): AdaptiveAvgPool2d(output_size=(7, 7))
  (flatten): Flatten(start_dim=1, end_dim=-1)
  (fc1): Linear(in_features=25088, out_features=4096, bias=True)
  (drop1): Dropout(p=0.5, inplace=False)
  (fc2): Linear(in_features=4096, out_features=1024, bias=True)
  (drop2): Dropout(p=0.5, inplace=False)
  (fc3): Linear(in_features=1024, out_features=10, bias=True)
)

# save the initialization of weights
torch.save(wild_model.state_dict(), 'wild_model_init.pt')

loss_object = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(wild_model.parameters())

def timelined_benchmark_wild(model, train_loader, test_loader, epochs):
    steps_acc, times_acc, values_acc = [], [], []

    start_time = time.perf_counter()
    print("start time: ", start_time)
    for epoch in range(epochs):
        model.train()
        print("training:")
        for steps, times, values, image, label in tqdm(train_loader, total=math.floor(len(img_path_train)/BATCH_SIZE)):
            for s in steps: steps_acc.extend(s)
            for t in times: times_acc.extend(t)
            for v in values: values_acc.extend(v)

            image, label = image.to(device), label.to(device)
            train_enter = time.perf_counter()
            optimizer.zero_grad()
            loss = loss_object(model(image), label)
            loss.backward()
            optimizer.step()
            train_elapsed = time.perf_counter() - train_enter

            steps_acc.append('Train')
            times_acc.append((train_enter, train_elapsed))
            values_acc.append((-1, epoch, -1))

        model.eval()
        print("testing:")
        with torch.no_grad():
            for steps, times, values, image, label in tqdm(test_loader, total=math.floor(len(img_path_test)/BATCH_SIZE)):
                for s in steps: steps_acc.extend(s)
                for t in times: times_acc.extend(t)
                for v in values: values_acc.extend(v)

                image, label = image.to(device), label.to(device)
                test_enter = time.perf_counter()
                model(image)
                test_elapsed = time.perf_counter() - test_enter
                steps_acc.append('Test')
                times_acc.append((test_enter, test_elapsed))
                values_acc.append((-1, epoch, -1))

    print("Execution time:", time.perf_counter() - start_time)
    return {'steps': np.array(steps_acc), 'times': np.array(times_acc, dtype=np.float64),
            'values': np.array(values_acc, dtype=np.int64)}

timeline_naive = timelined_benchmark_wild(wild_model, train_loader, test_loader, epochs=3)

start time:  1370371.388131665
training:

180it [01:54,  1.57it/s]                                                       

testing:

45it [00:27,  1.64it/s]                                                        

training:

180it [01:53,  1.58it/s]                                                       

testing:

45it [00:27,  1.62it/s]                                                        

training:

180it [01:54,  1.58it/s]                                                       

testing:

45it [00:27,  1.65it/s]                                                        

Execution time: 425.3732925138902

draw_timeline(timeline=timeline_naive, title="Naive", min_width=1)

Now try some data augmentation (transformation) to observe whether the accuracy and execution time increase or decrease.

## TODO: build `train_loader_augmentation` / `test_loader_augmentation`
## (data augmentation only -- NO optimization yet)
##
## Step 1. Define the transforms with transforms.Compose([...])
##   - augment_transform (training): add some random augmentation, e.g.
##     RandomCrop, RandomHorizontalFlip, ColorJitter, ... The output must still
##     be (3, IMAGE_SIZE_CROPPED, IMAGE_SIZE_CROPPED).
##   - test_transform (testing): NO random augmentation, only deterministic ops
##     (e.g. CenterCrop), so the test set is identical every epoch.
## Step 2. Build the datasets with TimeMeasuredDataset (see
##     dataset_generator_fun_train above), passing your transform via `transform=`.
## Step 3. Wrap them in DataLoaders:
##   - keep `batch_size=BATCH_SIZE` and `collate_fn=timing_collate`
##     (timelined_benchmark_wild() relies on it)
##   - use num_workers=0; do NOT set pin_memory / prefetch_factor yet
## Tip: plot a few samples from dataset_train_aug (like we did above) to check
##      that the augmentation looks the way you intended.

# augment_transform = transforms.Compose([...])
# test_transform    = transforms.Compose([...])
# dataset_train_aug = TimeMeasuredDataset(open_csv, read_image_file, 'train.csv', transform=...)
# dataset_test_aug  = TimeMeasuredDataset(open_csv, read_image_file, 'test.csv',  transform=...)
# train_loader_augmentation = DataLoader(dataset_train_aug, ...)
# test_loader_augmentation  = DataLoader(dataset_test_aug,  ...)

# load the same initialization of weights and re-train with the augmented pipeline
wild_model.load_state_dict(torch.load('wild_model_init.pt'))
timeline_augmentation = timelined_benchmark_wild(wild_model, train_loader_augmentation, test_loader_augmentation, epochs=3)

start time:  1371218.909736579
training:

180it [02:31,  1.19it/s]                                                       

testing:

45it [00:27,  1.61it/s]                                                        

training:

180it [02:29,  1.21it/s]                                                       

testing:

45it [00:27,  1.65it/s]                                                        

training:

180it [02:29,  1.20it/s]                                                       

testing:

45it [00:27,  1.64it/s]                                                        

Execution time: 532.3011030438356

draw_timeline(timeline=timeline_augmentation, title="Augmentation", min_width=1)

After trying data augmentation (transformation), it's time to optimize what you did above for better efficiency.

## TODO: build `train_loader_optimized` / `test_loader_optimized`
## (same transforms as above + DataLoader optimizations)
##
## Step 1. Create NEW dataset objects (don't reuse dataset_train_aug /
##     dataset_test_aug: TimeMeasuredDataset counts its calls internally, so a
##     reused dataset would mislabel the epochs in the timeline).
##     You can reuse augment_transform / test_transform from above.
##   (Optional) Caching: copy CachedTimeMeasuredDataset from the CIFAR-10
##     section into this part and use it instead. It caches the image BEFORE
##     the random augmentation, so augmentation still changes every epoch.
## Step 2. Wrap them in DataLoaders, keeping batch_size=BATCH_SIZE and
##     collate_fn=timing_collate, and add some of:
##   - num_workers=...          (e.g. the number of CPU cores)
##   - pin_memory=True          (only helps when training on a CUDA GPU)
##   - prefetch_factor=...      (requires num_workers > 0)
##   - persistent_workers=True  (required for caching to survive across epochs)

# dataset_train_opt = ...
# dataset_test_opt  = ...
# train_loader_optimized = DataLoader(dataset_train_opt, ...)
# test_loader_optimized  = DataLoader(dataset_test_opt,  ...)

# load the same initialization of weights and re-train with the optimized pipeline
wild_model.load_state_dict(torch.load('wild_model_init.pt'))
timeline_optimized = timelined_benchmark_wild(wild_model, train_loader_optimized, test_loader_optimized, epochs=3)

start time:  1371799.365239256
training:

180it [00:46,  3.85it/s]                                                       

testing:

45it [00:09,  4.97it/s]                                                        

training:

180it [00:47,  3.81it/s]                                                       

testing:

45it [00:09,  4.86it/s]                                                        

training:

180it [00:45,  3.94it/s]                                                       

testing:

45it [00:08,  5.01it/s]                                                        

Execution time: 166.99199603195302

draw_timeline(timeline=timeline_optimized, title="Optimized", min_width=1)