Introduction to Deep Learning with PyTorch#

This workshop introduces the essential concepts in deep learning and walks through the common steps in a deep learning workflow from data loading and preprocessing to training and model evaluation. Throughout the workshop, students participate in writing and executing simple deep learning programs using Pytorch – a popular Python library for developing, training, and deploying deep learning models.

Our final goal is an MNIST image-classification example: import image data, define a convolutional neural network (CNN), train it, and evaluate its predictions.

Before reaching that example, we will learn just enough about tensors, devices, autograd, and custom modules to understand what each part of the final workflow does.

0. Setup#

This workshop is designed for a Palmetto Jupyter session with a CUDA-enabled GPU. The notebook can fall back to CPU, but training will be slower.

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

# Fix the random seed so demonstrations are more reproducible.
torch.manual_seed(355)

# Use the first CUDA GPU when available; otherwise fall back to the CPU.
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

print("PyTorch version:", torch.__version__)
print("Device:", device)
if torch.cuda.is_available():
    print("GPU:", torch.cuda.get_device_name(0))
else:
    print("WARNING: No CUDA GPU was detected. Start a GPU-backed Palmetto session for this workshop.")

Preview: where we are going#

By the end, we will understand this repeated training pattern:

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

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

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

For now, simply notice the data, model, loss, gradients, and parameter update. We will return to each piece.

1. PyTorch basics: tensors#

PyTorch is a library for creating and training neural networks in Python. Its basic data structure is the tensor.

Tensors are similar to NumPy arrays, but they can run on GPUs and are optimized for automatic differentiation. We use them to store model inputs, outputs, parameters, and gradients.

Initializing a tensor#

# A range of values
print(torch.arange(10))
print(torch.arange(5, 15, 2))

# Directly from data
data = [[1, 2], [3, 4]]
# torch.tensor converts the nested Python list into a tensor.
x_data = torch.tensor(data)
print(x_data)

# Random and constant values
# The tuple defines two rows and three columns.
shape = (2, 3)
rand_tensor = torch.rand(shape)
ones_tensor = torch.ones(shape)
zeros_tensor = torch.zeros(shape)

print("Random Tensor:\n", rand_tensor)
print("Ones Tensor:\n", ones_tensor)
print("Zeros Tensor:\n", zeros_tensor)

# Tensors can be created from NumPy arrays
np_arr = np.array([[1, 2], [3, 4]])
tensor = torch.from_numpy(np_arr)
np_arr_2 = tensor.numpy()

print("Numpy array:\n", np_arr)
print("PyTorch tensor:\n", tensor)
print("Numpy array 2:\n", np_arr_2)

Tensor attributes#

tensor = torch.rand(3, 4)

print("Shape of tensor:", tensor.shape)
print("Datatype of tensor:", tensor.dtype)
print("Device tensor is stored on:", tensor.device)

Indexing and changing shape#

# Changing shape with view and reshape

# view uses the same storage from original tensor
x = torch.arange(12).view(3, 4) 
print("x view:",x) 
# reshape makes a copy of the tensor if needed (when tensor is not contiguous in memory)
x = torch.arange(12).reshape(3, 4)
print("x reshape:",x)

# A colon selects every value along that dimension.
print("Second column:", x[:, 1])
print("First row:", x[0])
print("Middle two rows:\n", x[1:3])
# flatten collapses all dimensions into one dimension.
print("Flattened:", x.flatten())

Image data sometimes arrives as (batch, height, width, channels), while PyTorch image models usually expect (batch, channels, height, width). permute rearranges the dimensions.

images = torch.rand(8, 28, 28, 3)
print("Before:", images.shape)

# Rearrange BHWC data into the BCHW order expected by PyTorch image models.
images = images.permute(0, 3, 1, 2)
print("After: ", images.shape)

Arithmetic and aggregation#

tensor = torch.arange(12, dtype=torch.float32).reshape(3, 4)
print(tensor)

# @ performs matrix multiplication; * performs element-wise multiplication.
print("Matrix multiplication:\n", tensor @ tensor.T)
print("Element-wise multiplication:\n", tensor * tensor)
print("Mean:", tensor.mean())
# axis=0 reduces the row dimension and returns one sum per column.
print("Sum along the rows:", tensor.sum(axis=0))
# axis=1 reduces the col dimension and returns one sum per row.
print("Sum along the cols:", tensor.sum(axis=1))

Quick check#

For a batch of grayscale images with shape (32, 1, 28, 28), what does each dimension represent?

2. GPU support and device selection#

GPUs can perform many similar operations in parallel, making them well suited to the large matrix operations used in neural networks.

By default, tensors are created on the CPU. A common PyTorch pattern is to define a device object and move tensors and models to it with .to(device). Tensors used in the same operation must be on the same device.

device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
print("Device:", device)

# Tensors are created on the CPU unless another device is specified.
x = torch.randn(3, 3)
print("Before:", x.device)

# .to(device) returns a tensor stored on the selected device.
x = x.to(device)
print("After: ", x.device)

3. Computation graphs and backpropagation#

One of the main reasons for using PyTorch is automatic differentiation. When a tensor has requires_grad=True, PyTorch records operations involving that tensor. Calling backward() computes the gradients needed to train a neural network.

Consider:

\[ y = \frac{1}{n}\sum_{i=1}^{n}\left[(x_i + 2)^2 + 3\right], \]

The gradient of \(y\) with respect to each element \(x_i\) is:

\[ \frac{\partial y}{\partial x_i} = \frac{2(x_i+2)}{n}. \]

For \(x=[0,1,2]\) and \(n=3\):

\[ x.\mathrm{grad} = \left[ \frac{4}{3}, 2, \frac{8}{3} \right] \approx [1.3333,\ 2.0000,\ 2.6667]. \]
x = torch.arange(3, dtype=torch.float32, requires_grad=True)
print("X:", x)

# Each operation becomes part of the computation graph.
a = x + 2
b = a ** 2
c = b + 3
y = c.mean()
print("Y:", y)

# Traverse the graph backward and store dy/dx in x.grad.
y.backward()
print("Gradient:", x.grad)

x.grad now contains \(\partial y / \partial x\). In a neural network, the same process computes a gradient for every learnable parameter.

4. Building a network with nn.Module#

The torch.nn package contains neural-network building blocks such as nn.Linear, nn.Conv2d, and activation functions.

To create a custom module, subclass nn.Module, define the layers in __init__, and describe the computation in forward.

class NeuralNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        # Flatten converts each 28 x 28 image into a vector of 784 values.
        self.flatten = nn.Flatten()
        # Sequential applies these layers in the listed order.
        self.layers = nn.Sequential(
            nn.Linear(28 * 28, 128),
            nn.ReLU(),
            nn.Linear(128, 10),
        )

    def forward(self, x):
        # Define how inputs flow through the model to produce logits.
        x = self.flatten(x)
        return self.layers(x)


# Move the model and its example input to the same device.
example_model = NeuralNetwork().to(device)
example_images = torch.randn(5, 1, 28, 28).to(device)
example_outputs = example_model(example_images)

print(example_model)
print("Input shape:", example_images.shape)
print("Output shape:", example_outputs.shape)

The output shape is (5, 10): one row for each image and one score for each digit class. These raw class scores are called logits.

For multiclass classification, nn.CrossEntropyLoss expects logits. Apply softmax only when probabilities are needed for interpretation—not inside the model during training.

5. PyTorch Quickstart: MNIST image classification#

We now have the pieces needed to understand a complete image-classification workflow. We will:

  1. import and inspect MNIST;

  2. define a CNN using a custom module;

  3. choose a loss function and optimizer;

  4. train the model using backpropagation;

  5. evaluate it on the test set.

Import the data#

Dataset represents individual samples. DataLoader groups those samples into mini-batches and yields one batch at a time. We shuffle the training data but not the test data.

BATCH_SIZE = 32

# Convert each image into a PyTorch tensor when it is loaded.
transform = transforms.Compose([
    transforms.ToTensor()
])

# download=True downloads MNIST only when it is not already in ./data.
trainset = torchvision.datasets.MNIST(
    root="./data",
    train=True,
    download=True,
    transform=transform,
)
trainloader = torch.utils.data.DataLoader(
    trainset,
    batch_size=BATCH_SIZE,
    # Shuffle training examples so batches change from epoch to epoch.
    shuffle=True,
    # Use worker processes to load batches while the GPU is computing.
    num_workers=8,
)

testset = torchvision.datasets.MNIST(
    root="./data",
    train=False,
    download=True,
    transform=transform,
)
testloader = torch.utils.data.DataLoader(
    testset,
    batch_size=BATCH_SIZE,
    # Test order does not affect the evaluation result.
    shuffle=False,
    num_workers=8,
)

Explore the data#

dataiter = iter(trainloader)
# Request one mini-batch from the data loader.
images, labels = next(dataiter)

print("Image batch dimensions:", images.shape)
print("Image label dimensions:", labels.shape)
print("First labels:", labels[:10].tolist())

grid = torchvision.utils.make_grid(images[:16], nrow=8)
# Matplotlib expects height, width, channels rather than channels first.
npimg = grid.numpy()

plt.figure(figsize=(12, 3))
plt.imshow(np.transpose(npimg, (1, 2, 0)), cmap="gray")
plt.axis("off")
plt.show()

The model#

The model uses a compact architecture with one convolutional layer and two linear layers:

  • one convolutional layer extracts 32 feature maps;

  • ReLU adds nonlinearity;

  • flatten converts feature maps into a vector;

  • two linear layers produce 10 class logits.

A 3×3 convolution changes a 28×28 image into 26×26 feature maps, so the first linear layer receives 32 * 26 * 26 values per image.

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

        # Learn 32 different 3 x 3 filters from the grayscale input.
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3)
        # The convolution produces 32 feature maps, each 26 x 26.
        self.d1 = nn.Linear(32 * 26 * 26, 128)
        # Produce one score for each digit class, 0 through 9.
        self.d2 = nn.Linear(128, 10)

    def forward(self, x):
        # Extract local image features and apply a nonlinearity.
        x = self.conv1(x)
        x = F.relu(x)

        # Preserve the batch dimension while flattening each image.
        x = x.flatten(start_dim=1)

        x = self.d1(x)
        x = F.relu(x)

        # Return raw class scores for CrossEntropyLoss.
        logits = self.d2(x)
        return logits

The model returns logits because nn.CrossEntropyLoss expects logits and applies the necessary normalization internally. Use softmax only when probabilities are needed for interpretation.

model = MyModel().to(device)

# Test the model with one batch before training.
with torch.no_grad():
    # Images and model parameters must be on the same device.
    out = model(images.to(device))

print("Input shape:", images.shape)
print("Output shape:", out.shape)
print("Trainable parameters:", f"{sum(p.numel() for p in model.parameters()):,}")

Before training, the model’s parameters are random. Convert its logits to probabilities only for display.

with torch.no_grad():
    # Softmax converts the ten logits into probabilities that sum to one.
    probabilities = F.softmax(out, dim=1).cpu()

fig, axes = plt.subplots(2, 8, figsize=(14, 4))
for i in range(8):
    # Show the image above its probability distribution.
    axes[0, i].imshow(images[i].squeeze(), cmap="gray")
    axes[0, i].set_title(f"true: {labels[i].item()}")
    axes[0, i].axis("off")

    axes[1, i].bar(np.arange(10), probabilities[i].numpy())
    axes[1, i].set_xticks(np.arange(10))
    axes[1, i].set_ylim(0, 1)

plt.tight_layout()
plt.show()

Predict before training#

What test accuracy do you think this model will achieve after five epochs?

Train the model#

We use cross-entropy loss and the Adam optimizer. Each mini-batch performs a forward pass, measures the loss, computes gradients, and updates the parameters.

learning_rate = 0.001
num_epochs = 5

# Start with a fresh, randomly initialized model for training.
model = MyModel().to(device)
# CrossEntropyLoss compares class logits with integer class labels.
criterion = nn.CrossEntropyLoss()
# Adam updates every learnable parameter returned by model.parameters().
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

for epoch in range(num_epochs):
    # Track totals across every mini-batch in this epoch.
    train_running_loss = 0.0
    train_correct = 0
    train_examples = 0

    # Enable training behavior for all layers in the model.
    model.train()

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

        # Forward pass, loss, and backpropagation
        outs = model(images)
        loss = criterion(outs, labels)
        # Gradients accumulate by default, so clear the previous batch first.
        optimizer.zero_grad()
        loss.backward()

        # Update model parameters
        optimizer.step()

        # Multiply by batch size so the final loss is an example-weighted mean.
        train_running_loss += loss.detach().item() * labels.size(0)
        # The largest logit is the predicted class.
        train_correct += (outs.argmax(dim=1) == labels).sum().item()
        train_examples += labels.size(0)

    print(
        f"Epoch: {epoch + 1} | "
        f"Loss: {train_running_loss / train_examples:.4f} | "
        f"Train Accuracy: {100 * train_correct / train_examples:.2f}%"
    )

Evaluate the model#

test_correct = 0
test_examples = 0

# Switch layers to evaluation behavior and stop recording gradients.
model.eval()
with torch.no_grad():
    for images, labels in testloader:
        images = images.to(device)
        labels = labels.to(device)

        outputs = model(images)
        # Count correct predictions across all test batches.
        test_correct += (outputs.argmax(dim=1) == labels).sum().item()
        test_examples += labels.size(0)

print(f"Test Accuracy: {100 * test_correct / test_examples:.2f}%")

# Look at predictions from one test batch.
images, labels = next(iter(testloader))

model.eval()
with torch.no_grad():
    outputs = model(images.to(device))
    # Move probabilities back to CPU for plotting with Matplotlib.
    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: {labels[i].item()}")
    axes[0, i].axis("off")

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

plt.tight_layout()
plt.show()

Coding challenge#

Modify the model to use two convolutional layers. Use 16 output channels in the first convolution, choose the number of channels in the second, and make sure the final output contains 10 logits.

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

        # HINT: Start with MyModel. Add a second convolutional layer
        # between conv1 and d1. The first layer should output 16 channels.
        # What must the second layer's in_channels be?
        # INSERT YOUR LAYERS HERE
        
        
        
    def forward(self, x):
        # HINT: Follow MyModel's forward pass, applying ReLU after
        # each convolution. Flatten before the linear layers.
        # INSERT YOUR CODE HERE
        
        
        return x

Model Check#

model2 = MyModel2()
with torch.no_grad():
    print("Output shape:", model2(images.to(device)).shape) # Should be (batch_size, 10)
Show one possible solution 
class MyModel2(nn.Module):
    def __init__(self):
        super(MyModel2, self).__init__()

        # First convolution: (1, 28, 28) → (16, 26, 26)
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=16, kernel_size=3)

        # Second convolution: (16, 26, 26) → (32, 24, 24)
        self.conv2 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3)

        # Flattened feature maps → 128 values → 10 digit scores
        self.d1 = nn.Linear(32 * 24 * 24, 128)
        self.d2 = nn.Linear(128, 10)

    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)

Final words#

You can now explain the main building blocks of a PyTorch image-classification workflow:

  • tensors represent the data and parameters;

  • .to(device) moves work to the GPU;

  • autograd computes gradients;

  • nn.Module organizes a model;

  • Dataset and DataLoader provide mini-batches;

  • the training loop performs forward passes, backpropagation, and parameter updates;

  • evaluation measures performance on unseen data.

The same overall structure scales from this basic CNN to much larger deep-learning models.