PyTorch Convolutional Neural Network

PyTorch Convolutional Neural Network (CNN) is a class of deep learning models specifically designed to process data with a grid-like topology (such as images).

CNN is a core technology for computer vision tasks (such as image classification, object detection, and segmentation).

The following diagram shows the structure and workflow of a typical Convolutional Neural Network (CNN) used for image recognition tasks.

In the figure, the CNN output layer gives probabilities for three classes: Donald (0.2), Goofy (0.1), and Tweety (0.7), indicating that the network believes the input image is most likely Tweety.

The following is a brief description of each part:

  • Input Image: The raw image data received by the network.

  • Convolution: A kernel slides over the input image to extract features and generate feature maps.

  • Pooling: Usually applied after the convolutional layer, max pooling or average pooling reduces the size of feature maps while preserving important features, producing pooled feature maps.

  • Feature Extraction: The combination of multiple convolutional and pooling layers gradually extracts high-level features from the image.

  • Flatten Layer: Converts the multi-dimensional feature maps into a one-dimensional vector for input into the fully connected layer.

  • Fully Connected Layer: Similar to traditional neural network layers, it maps the extracted features to output classes.

  • Classification: The output layer of the network performs classification based on the output of the fully connected layer.

  • Probabilistic Distribution: The output layer gives the probability for each class, representing the likelihood that the input image belongs to each class.

Basic Structure of Convolutional Neural Networks

1. Input Layer

It receives the raw image data. An image is typically represented as a three-dimensional array, where two dimensions represent the width and height of the image, and the third dimension represents the color channels (for example, an RGB image has three channels).

2. Convolutional Layer

It uses convolution kernels to extract local features, such as edges and textures.

Formula:

  • x: input image.
  • k: convolution kernel (weight matrix).
  • b: bias.

A group of learnable filters (or convolution kernels) is applied to the input image via convolution operations to extract local features.

Each filter slides over the input image and generates a feature map, representing the filter's activation at different positions.

A convolutional layer can have multiple filters; each filter generates a feature map, and all feature maps form a set of feature maps.

3. Activation Function

A nonlinear activation function, such as ReLU (Rectified Linear Unit), is usually applied after the convolutional layer to introduce nonlinearity, enabling the network to learn more complex patterns.

The ReLU function is defined as:f(x)=max(0,x), that is, if the input is less than 0, output 0; otherwise, output the input value.

4. Pooling Layer

  • It is used to reduce the spatial dimensions of feature maps, decreasing the computational cost and number of parameters while retaining the most important feature information.
  • The most common pooling operations are max pooling and average pooling.
  • Max pooling selects the maximum value in the region, while average pooling computes the average value in the region.

5. Normalization Layer (optional)

  • For example, Local Response Normalization (LRN) or Batch Normalization.
  • These layers help accelerate the training process and improve model stability.

6. Fully Connected Layer

  • At the end of a CNN, the feature maps extracted by earlier layers are flattened into a one-dimensional vector and then fed into the fully connected layer.
  • Every neuron in the fully connected layer is connected to all neurons in the previous layer, combining features for the final classification or regression.

7. Output Layer

Depending on the task, the output layer can take different forms.

For classification tasks, the Softmax function is typically used to convert the output into a probability distribution, representing the probability of the input belonging to each class.

8. Loss Function

It is used to measure the difference between the model's predictions and the true labels.

Common loss functions include cross-entropy loss for multi-class classification tasks and mean squared error (MSE) for regression tasks.

9. Optimizer

It is used to update the network's weights based on the gradient of the loss function. Common optimizers include stochastic gradient descent (SGD), Adam, RMSprop, etc.

10. Regularization (optional)

It includes techniques such as Dropout and L1/L2 regularization to prevent the model from overfitting.

These layers can be stacked to form deeper network structures to improve the model's learning capability.

The depth and complexity of a CNN can be adjusted according to the requirements of the task.


PyTorch Implementation of a CNN Example

The following example shows how to build a simple CNN model with PyTorch for digit classification on the MNIST dataset.

Main Steps:

  • Data Loading and Preprocessing: Usetorchvisionload and preprocess the MNIST data.
  • Model Construction: Define convolutional, pooling, and fully connected layers.
  • Training: Train the model using a loss function and an optimizer.
  • Evaluation: Calculate the model's accuracy on the test set.
  • Visualization: Display some test samples and their prediction results.

1. Import Required Libraries

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

2. Data Loading

Use the MNIST dataset provided by torchvision to load and preprocess the data.

Example

transform = transforms.Compose([
    transforms.ToTensor(),  # Convert to tensor
    transforms.Normalize((0.5,), (0.5,))  # Normalize to [-1, 1]
])

# Load MNIST dataset
train_dataset = datasets.MNIST(root='./data', train=True, transform=transform, download=True)
test_dataset = datasets.MNIST(root='./data', train=False, transform=transform, download=True)

train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=64, shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=64, shuffle=False)

3. Define the CNN Model

Usenn.Modulebuild a CNN.

Example

class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        # Define convolutional layer: 1 input channel, 32 output channels, kernel size 3x3
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1)
        # Define convolutional layer: 32 input channels, 64 output channels
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
        # Define fully connected layer
        self.fc1 = nn.Linear(64 * 7 * 7, 128)  # Input size = feature map size * number of channels
        self.fc2 = nn.Linear(128, 10)  # 10 classes

    def forward(self, x):
        x = F.relu(self.conv1(x))  # First convolutional layer + ReLU
        x = F.max_pool2d(x, 2)     # Max pooling
        x = F.relu(self.conv2(x))  # Second convolutional layer + ReLU
        x = F.max_pool2d(x, 2)     # Max pooling
        x = x.view(-1, 64 * 7 * 7) # Flatten operation
        x = F.relu(self.fc1(x))    # Fully connected layer + ReLU
        x = self.fc2(x)            # Fully connected layer output
        return x

# Create model instance
model = SimpleCNN()

4. Define Loss Function and Optimizer

Use cross-entropy loss and stochastic gradient descent optimizer.

criterion = nn.CrossEntropyLoss()  # 多分类交叉熵损失
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)  # 学习率和动量

5. Train the Model

Train the model for 5 epochs, and output the training loss after each epoch.

Example

num_epochs = 5
model.train()  # Set to training mode

for epoch in range(num_epochs):
    total_loss = 0
    for images, labels in train_loader:
        # Forward pass
        outputs = model(images)
        loss = criterion(outputs, labels)
       
        # Backward pass
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        total_loss += loss.item()
    print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {total_loss / len(train_loader):.4f}")

6. Test the Model

Evaluate the model's accuracy on the test set.

Example

model.eval()  # Set to evaluation mode
correct = 0
total = 0

with torch.no_grad():  # No gradient calculation needed during evaluation
    for images, labels in test_loader:
        outputs = model(images)
        _, predicted = torch.max(outputs, 1)  # Predict classes
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

accuracy = 100 * correct / total
print(f"Test Accuracy: {accuracy:.2f}%")

7. Complete Code

The complete code is as follows:

Example

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms


# 1. Data loading and preprocessing
transform = transforms.Compose([
    transforms.ToTensor(),  # Convert to tensor
    transforms.Normalize((0.5,), (0.5,))  # Normalize to [-1, 1]
])

# Load MNIST dataset
train_dataset = datasets.MNIST(root='./data', train=True, transform=transform, download=True)
test_dataset = datasets.MNIST(root='./data', train=False, transform=transform, download=True)

train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=64, shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=64, shuffle=False)

# 2. Define CNN model
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        # Define convolutional layers
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1)  # Input 1 channel, output 32 channels
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)  # Input 32 channels, output 64 channels
        # Define fully connected layers
        self.fc1 = nn.Linear(64 * 7 * 7, 128)  # Flatten and feed into fully connected layer
        self.fc2 = nn.Linear(128, 10)  # 10 classes

    def forward(self, x):
        x = F.relu(self.conv1(x))  # First convolution + ReLU
        x = F.max_pool2d(x, 2)     # Max pooling
        x = F.relu(self.conv2(x))  # Second convolution + ReLU
        x = F.max_pool2d(x, 2)     # Max pooling
        x = x.view(-1, 64 * 7 * 7) # Flatten
        x = F.relu(self.fc1(x))    # Fully connected layer + ReLU
        x = self.fc2(x)            # Final output layer
        return x

# Create model instance
model = SimpleCNN()

# 3. Define loss function and optimizer
criterion = nn.CrossEntropyLoss()  # Multi-class cross-entropy loss
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

# 4. Model training
num_epochs = 5
model.train()  # Set model to training mode

for epoch in range(num_epochs):
    total_loss = 0
    for images, labels in train_loader:
        outputs = model(images)  # Forward pass
        loss = criterion(outputs, labels)  # Compute loss

        optimizer.zero_grad()  # Clear gradients
        loss.backward()  # Backward pass
        optimizer.step()  # Update parameters

        total_loss += loss.item()

    print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {total_loss / len(train_loader):.4f}")

# 5. Model testing
model.eval()  # Set model to evaluation mode
correct = 0
total = 0

with torch.no_grad():  # Disable gradient computation
    for images, labels in test_loader:
        outputs = model(images)
        _, predicted = torch.max(outputs, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

accuracy = 100 * correct / total
print(f"Test Accuracy: {accuracy:.2f}%")

Explanation of Running Results

1. Output training loss

The code outputs an average loss after each epoch, for example:

Epoch [1/5], Loss: 0.2325
Epoch [2/5], Loss: 0.0526
Epoch [3/5], Loss: 0.0366
Epoch [4/5], Loss: 0.0273
Epoch [5/5], Loss: 0.0221

Explanation:The gradual decrease in loss indicates that the model is converging step by step.

2. Accuracy on the test set

The code outputs the final classification accuracy on the test set, for example:

Test Accuracy: 98.96%

Explanation:The model achieves a classification accuracy of 98.96% on the MNIST test set, which is a good result for a simple CNN model.

7. Visualize Results

We can visualize some samples and their prediction results in the test data.

Example

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
import matplotlib.pyplot as plt

# 1. Data loading and preprocessing
transform = transforms.Compose([
    transforms.ToTensor(),  # Convert to tensor
    transforms.Normalize((0.5,), (0.5,))  # Normalize to [-1, 1]
])

# Load MNIST dataset
train_dataset = datasets.MNIST(root='./data', train=True, transform=transform, download=True)
test_dataset = datasets.MNIST(root='./data', train=False, transform=transform, download=True)

train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=64, shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=64, shuffle=False)

# 2. Define CNN model
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        # Define convolutional layers
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1)  # Input 1 channel, output 32 channels
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)  # Input 32 channels, output 64 channels
        # Define fully connected layers
        self.fc1 = nn.Linear(64 * 7 * 7, 128)  # Flatten and feed into fully connected layer
        self.fc2 = nn.Linear(128, 10)  # 10 classes

    def forward(self, x):
        x = F.relu(self.conv1(x))  # First convolution + ReLU
        x = F.max_pool2d(x, 2)     # Max pooling
        x = F.relu(self.conv2(x))  # Second convolution + ReLU
        x = F.max_pool2d(x, 2)     # Max pooling
        x = x.view(-1, 64 * 7 * 7) # Flatten
        x = F.relu(self.fc1(x))    # Fully connected layer + ReLU
        x = self.fc2(x)            # Final output layer
        return x

# Create model instance
model = SimpleCNN()

# 3. Define loss function and optimizer
criterion = nn.CrossEntropyLoss()  # Multi-class cross-entropy loss
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

# 4. Model training
num_epochs = 5
model.train()  # Set model to training mode

for epoch in range(num_epochs):
    total_loss = 0
    for images, labels in train_loader:
        outputs = model(images)  # Forward pass
        loss = criterion(outputs, labels)  # Compute loss

        optimizer.zero_grad()  # Clear gradients
        loss.backward()  # Backward pass
        optimizer.step()  # Update parameters

        total_loss += loss.item()

    print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {total_loss / len(train_loader):.4f}")

# 5. Model testing
model.eval()  # Set model to evaluation mode
correct = 0
total = 0

with torch.no_grad():  # Disable gradient computation
    for images, labels in test_loader:
        outputs = model(images)
        _, predicted = torch.max(outputs, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

accuracy = 100 * correct / total
print(f"Test Accuracy: {accuracy:.2f}%")

# 6. Visualize test results
dataiter = iter(test_loader)
images, labels = next(dataiter)
outputs = model(images)
_, predictions = torch.max(outputs, 1)

fig, axes = plt.subplots(1, 6, figsize=(12, 4))
for i in range(6):
    axes[i].imshow(images[i][0], cmap='gray')
    axes[i].set_title(f"Label: {labels[i]}\nPred: {predictions[i]}")
    axes[i].axis('off')
plt.show()

The visualization results show the actual labels and predicted values of 6 test samples, for example:

  • The upper left corner of the image is the handwritten digit.

  • The title section shows the model's predicted value and the true label.

Note Several Common Errors

When downloading the MNIST dataset, the download fails due to SSL certificate verification failure.

Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz
Failed to download (trying next):
<urlopen error [SSL: CERTIFICATE_VERIFY_FAILED] certificate verify failed: unable to get local issuer certificate (_ssl.c:1129)>
...

Solutions

Open a terminal and run the following command to install/update the SSL certificate:

/Applications/Python\ 3.x/Install\ Certificates.command

Replace 3.x with the Python version you actually installed.

After running this command, re-run your Python code and it should be able to download the MNIST dataset successfully.

Other extensions