# Convolutional Neural Networks (CNNs) In Pytorch

#34 Jan 20, 2025 Architecture

https://www.aibutsimple.com/p/convolutional-neural-networks-cnns-in-pytorch

Convolutional Neural Networks (CNNs) are the foundational model of modern computer vision applications.

Specifically, they are a special type of neural network designed to process grid-like data, such as images. With their high performance and efficient architecture, they excel in computer vision tasks such as object detection and image classification.

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/fccfb0ba-76e7-4156-b32d-5270bab8941d/image.png)

CNNs are particularly effective for image classification tasks due to their ability to capture patterns in images (such as edges, textures, or shapes) through convolutional layers.

In this issue, we’ll use PyTorch to build CNNs, going over model architecture, essential layers, and best practices.

---

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/86a55dfa-d764-446c-ae07-6208a408676f/image.png)

### Fact-based news without bias awaits. Make 1440 your choice today.

Overwhelmed by biased news? Cut through the clutter and get straight facts with your daily 1440 digest. From politics to sports, join millions who start their day informed.

[Sign up now!](https://l.join1440.com/bh?utm_source=beehiiv&utm_medium=cpc&utm_campaign=ai_but_simple_XIAEVJHTGJ&utm_content=prospecting_humans&_bhiiv=opp_ec25a85c-414a-4093-ac5b-c804a095bb3f_1b75ca79)

---

Before continuing, we recommend having a good understanding of CNNs—especially how they work both conceptually and mathematically. If you don’t already, we have past issues explaining CNNs thoroughly below:

- [CNNs, Explained Mathematically (Intermediate)](https://www.aibutsimple.com/p/convolutional-neural-networks-explained-mathematically?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=convolutional-neural-networks-cnns-in-pytorch)
- [Convolutional Neural Networks, Simplified (Beginner)](https://www.aibutsimple.com/p/convolutional-neural-networks-simplified?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=convolutional-neural-networks-cnns-in-pytorch)

Additionally, a basic understanding of Python is needed to follow along, specifically some familiarity with the PyTorch framework.

The code in this issue can be executed on a normal home PC and does not require a large amount of VRAM. It can also be executed on hosted notebooks such as Google Colab.

### Worked Example: Simple CNN on CIFAR-10

---

To start, we’re going to create a simple CNN to tackle the [CIFAR-10 dataset](https://www.cs.toronto.edu/~kriz/cifar.html?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=convolutional-neural-networks-cnns-in-pytorch), a dataset of various 32×32 images belonging to 10 distinct classes. It is a classic multi-class classification problem.

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/2d27e40b-4f34-451e-9591-ab73dd0cf31d/image.png)

#### Imports

Let’s start by importing the necessary libraries. PyTorch is the main library we’ll be using for building and training the neural network.

We’ll also use torchvision to manage the dataset and apply transformations such as normalization. Tqdm is used here to display a progress bar during training.

```
import torch
import torchvision.datasets as datasets
import torchvision.transforms as transforms
import torch.nn.functional as F
from torch import nn
from torch.utils.data import DataLoader
from tqdm import tqdm
```

#### Set Up Device

PyTorch can run on both CPUs and GPUs. We’ll initialize a variable called “device” storing a string—we use a GPU (cuda) if it is available; otherwise, we’ll use a CPU.

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

#### Define CNN Architecture

For our model, we’ll create a simple CNN with three convolutional layers, then a fully connected layer to generate the predictions.

Our network will use 3 “standard” CNN blocks. A “standard” block contains a convolutional layer that maintains the resolution, followed by an activation function (usually ReLU), then a pooling layer to reduce the resolution.

- This block originates from VGGNet, a pioneer model at the time it was released.

The architecture of our model is shown below:

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/e07533aa-f158-4e11-835f-f6423c483254/Copy_of_Copy_of_Copy_of_Copy_of_Add_a_heading__720_x_600_px___1_.png)

We can define this model by writing a Python class (which we’ll name CNN) while using PyTorch’s neural network (nn) module:

```
class CNN(nn.Module):
    def __init__(self, in_channels=3, num_classes=10):
        """
        Define the layers of the convolutional neural network.
        """
        super(CNN, self).__init__()

        # 1st convolutional layer: 3 input channels (RGB), 32 output channels, 3x3 kernel, stride 1, padding 1 (Preserve Resolution)
        self.conv1 = nn.Conv2d(in_channels=in_channels, out_channels=32, kernel_size=3, stride=1, padding=1)

        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=1, padding=1)

        self.conv3 = nn.Conv2d(in_channels=64, 
out_channels=128, kernel_size=3, stride=1, padding=1)

        # Max pooling layer: 2x2 window, stride 2 (Halve Resolution)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)

        # Fully connected layer: 128*4*4 input features (flattened, after 3 pooling), 10 output features (num_classes)
        self.fc1 = nn.Linear(128 * 4 * 4, num_classes)

    def forward(self, x):
        """
        Define the forward pass of the neural network.
        """
        x = F.relu(self.conv1(x))  # Apply convolution and ReLU
        x = self.pool(x)           # Apply max pooling
        x = F.relu(self.conv2(x))
        x = self.pool(x)
        x = F.relu(self.conv3(x))
        x = self.pool(x)

        x = x.reshape(x.shape[0], -1)  # Flatten the tensor
        x = self.fc1(x)            # Apply FC layer

        return x
```

Above, the constructor for the CNN has 2 parameters:

- in_channels (int): The number of channels in the input image. CIFAR-10 images have 3 color channels (RGB), so the in_channels defaults to 3.
- num_classes (int): The number of classes in our final prediction. CIFAR-10 has a total of 10 different classifications (airplane, automobile, bird, etc.).

The number of input channels can change depending on the type of data used. For instance, an RGB image has 3 color channels while grayscale images only have 1.

In the convolutional layers themselves (Conv2d), out_channels refers to the number of output channels (or feature maps) produced by the layer.

When performing convolutions, each filter/kernel produces one feature map, and the number of filters determines the number of output channels. For instance, if a convolutional layer has 64 filters, then the number of out_channels would be 64.

Out_channels is specified in the code itself in arbitrary choice—if the task is more complex, you may need more out_channels and more convolutional layers to capture the patterns in the images.

##### Forward Propagation

Moving on, the forward method is used for the forward propagation of the CNN. The parameter x represents the input tensor of CIFAR-10 images.

def forward(self, x):

The input tensor (x) for a batch of images in PyTorch usually has dimensions of (N, C, H, W), where:

- N is the batch size.
- C is the number of channels (this corresponds to in_channels).
- H is the height of the image.
- W is the width of the image.

The input tensor passes through a convolutional layer, which produces an output tensor with dimensions (N, out_channels, H’, W’).

- H' and W' are the height and width of the output feature maps, which may change based on padding, stride, and kernel size.
- In our case, we use a 3×3 kernel with a stride of 1 and a padding of 1, which maintains the resolution of the feature maps (H and W remain unchanged).

self.conv1 = nn.Conv2d(in_channels=in_channels, out_channels=32, kernel_size=3, stride=1, padding=1)

---

If you don’t understand what stride or padding means, or are confused with any other terminology, please consult our past issue introducing the CNN:

- [Convolutional Neural Networks, Simplified (Beginner)](https://www.aibutsimple.com/p/convolutional-neural-networks-simplified?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=convolutional-neural-networks-cnns-in-pytorch)

---

We then pass the output tensor through a ReLU activation function to add nonlinearity and to follow the format of the “standard” CNN block. Another reason for this activation is to enhance feature extraction—it ensures only the most important patterns are preserved while filtering out less important signals.

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

At this point, we can use pooling to reduce the dimensionality, extract the most important features, and improve translational invariance (being able to classify accurately regardless of the position of the object).

self.pool = nn.MaxPool2d(kernel_size=2, stride=2)

In our CNN, we use max pooling, which takes the maximum value in a pooling window. This changes the height and width of the image but not the number of feature maps or channels.

- Specifically, we use a 2×2 pooling window with padding 2 and stride 2, which halves the image size (N/2, W/2).

So far, our forward pass process has followed the “standard” CNN block. If you look at the code, you’ll notice that the sequence aligns perfectly: The input tensor is passed through a convolution layer, then a ReLU activation function, then a max pooling layer.

This process (convolution, ReLU, pooling) is repeated three times before flattening the tensor into a 1D vector and feeding it through a fully connected layer with 10 neurons to predict the class of the image.

##### Input Tensor Dimension Changes

As a quick summary, here’s how the size of our input changes over time through the network:

- Input image (3x32×32)
- First convolutional layer (3x32×32) → (32×32×32) (32 out_channels as specified)
- ReLU activation function (no changes to dimensions)
- First max pooling layer (32×32×32) → (32x16×16) (Width, Height halved)
- Second convolutional layer (32x16×16) → (64x16×16) (64 out_channels as specified)
- ReLU activation function (no changes to dimensions)
- Second max pooling layer (64x16×16) → (64×8×8) (Width, Height halved)
- Third convolutional layer (64×8×8) → (128×8×8) (128 out_channels as specified)
- Third max pooling layer (128×8×8) → (128×4×4) (Width, Height halved)

At the end, the max pooling layer has dimensions 128×4×4, which is flattened to a size of 128×4×4=2048. This is passed as input to our fully connected layer, which has 10 neurons—1 for each class. The neuron with the highest numerical value will be the predicted class.

#### Define Hyperparameters

Hyperparameters are configuration settings used to tune how the model is trained. In our case, the hyperparameters are simply numerical values.

```
num_classes = 10  # 10 image classes
learning_rate = 0.001
batch_size = 64
num_epochs = 10  # Small for demonstration purposes
```

Above, we define variables for the number of classes (10), the learning rate (used for our optimizer, 0.001), and the batch size (64, chosen as a standard).

We also set the number of epochs in the training process to 10, a very short amount of time for demonstration purposes (you may increase this number to allow the model to fully learn the dataset).

#### Load The Dataset

We load the CIFAR-10 dataset through torchvision itself, using the transforms.compose method to standardize and normalize the data, finally reformatting it into a tensor.

```
# Use transforms.compose method to standardize, normalize, and reformat into usable condition
all_transforms = transforms.Compose([transforms.ToTensor(),
            transforms.Normalize(mean=[0.4914, 0.4822, 0.4465],
            std=[0.2023, 0.1994, 0.2010])])

# Create Training dataset
train_dataset = datasets.CIFAR10(root = './data',
     train = True, transform = all_transforms, download = True)

# Create Test dataset
test_dataset = datasets.CIFAR10(root = './data',
     train = False, transform = all_transforms, download=True)

# Create dataloaders for easier processing
train_loader = torch.utils.data.DataLoader(dataset =      train_dataset, batch_size = batch_size, shuffle = True)


test_loader = torch.utils.data.DataLoader(dataset =  test_dataset, batch_size = batch_size,shuffle = True)
```

The output should look something like this if you download the dataset for the first time:

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/a51dea13-932b-43ba-8a03-65f77d3f3301/image.png)

For our example, we use dataloaders that simplify and optimize the process of feeding data into the CNN. They are the mechanisms that process the dataset into smaller batches instead of loading the entire dataset into memory.

We can also shuffle the data using shuffle=True, introducing randomness to ensure the model does not remember the order of the data.

#### Setup

```
model = CNN(num_classes = num_classes).to(device)

# Set Loss function with to Cross Entropy Loss
criterion = nn.CrossEntropyLoss()

# Set optimizer to Adam
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
```

We create a CNN object and send it to our device, which in this case is the GPU. Then, we initialize a variable “criterion” to the cross-entropy loss since we are doing multiclass classification.

We’ll use Adam as our optimization algorithm to train the CNN—we use the learning rate set previously.

#### Train the Model

```
loss_history = []
for epoch in range(num_epochs):
    epoch_loss = 0.0
    print(f"Epoch [{epoch + 1}/{num_epochs}]")
    for batch_index, (data, targets) in enumerate(tqdm(train_loader)):
        # Move data and targets to the device (GPU/CPU)
        data = data.to(device)
        targets = targets.to(device)

        # Forward pass: compute the model output
        scores = model(data)
        loss = criterion(scores, targets)

        # Backward pass: compute the gradients
        optimizer.zero_grad()
        loss.backward()

        # Optimization step: update the model parameters
        optimizer.step()

        # Accumulate epoch loss
        epoch_loss += loss.item()

    loss_history.append(epoch_loss / len(train_loader))
```

We use two for loops, one looping each epoch and one looping over the entire dataset in batches. In the inner loop, we send the data and the labels to our GPU.

data = data.to(device)

targets = targets.to(device)

We then give the model our data to produce the predictions and calculate the loss for the batch by comparing the predicted class to the true label.

scores = model(data)

loss = criterion(scores, targets)

After computing the loss, we’ll compute the gradients via backpropagation, then update the model parameters (doing the learning).

optimizer.zero_grad()

loss.backward()

optimizer.step()

At the end of the inner loop, we accumulate the loss for the entire epoch by continuously adding the batch loss.

Then, in the outer loop, we divide the total epoch loss by the number of batches to obtain the average loss for the epoch. We add this to a list, storing the loss for each epoch for later use.

Now, we can finally train the model! We’ll train it for 10 epochs:

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/1afe62b5-5518-4e83-9f81-2f549c19430b/image.png)

This entire process took around 2-3 minutes on Google Colab. Depending on the type of GPU you have, it will usually run a lot faster.

#### Plot Loss Curve

```
import matplotlib.pyplot as plt
plt.figure(figsize=(11, 7)) # set the figure size

# plot the loss history
plt.plot(range(1, num_epochs + 1), loss_history, marker='o', linestyle='-', label='Training Loss') 

plt.title("Loss Curve")
plt.xlabel("Epoch")
plt.ylabel("Loss")
plt.xticks(range(1, num_epochs + 1))
plt.legend()
plt.grid()
plt.show()
```

Using the list of losses, we can plot it to visualize our model’s training. Here’s what the loss curve looks like after 10 epochs:

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/d21bdc90-7614-470a-b397-668a1d6cf6e2/image.png)

We observe that the model is learning very steadily, but it is important to keep in mind that we have only trained the CNN for 10 epochs, which is a very small number.

#### Check Model Accuracy

At last, we should check the model’s accuracy to determine its performance on both the train and test datasets. We’ll define a method to check it:

```
def check_accuracy(loader, model):
    """
    Checks the accuracy of a model on a given Dataloader.
    """
    if loader.dataset.train:
        print("Checking Model Accuracy on Training Data:")
    else:
        print("Checking Model Accuracy on Test Data:")

    correct = 0 # Number of correct predictions
    samples = 0 # Number of total samples
    model.eval()  # Set the model to evaluation mode

    with torch.no_grad():  # Disable gradient calculation
        for x, y in loader:
            x = x.to(device)
            y = y.to(device)

            # Forward pass: compute the model output
            scores = model(x)
            # Get the index of max probability
            _, predictions = scores.max(1)
            # Accumulate correct predictions
            correct += (predictions == y).sum()
            # Accumulate total samples
            samples += predictions.size(0)

        # Calculate accuracy
        accuracy = float(correct) / float(samples) * 100
        print(f"Correct Predictions: {correct}/{samples}")
        print(f"Accuracy: {accuracy:.2f}%")

    model.train()  # Set the model back to training mode

# Check accuracy on both train and test sets
check_accuracy(train_loader, model)
check_accuracy(test_loader, model)
```

Above, we disable gradient calculation to reduce memory usage and speed up computations during the prediction process (inference).

with torch.no_grad()

Running the entire code gives us:

Checking Model Accuracy on Training Data:

Correct Predictions: 46338/50000

Accuracy: 92.68%

Checking Model Accuracy on Test Data:

Correct Predictions: 7480/10000

Accuracy: 74.80%

Although the test accuracy needs a lot of improvement, it isn’t too bad for 10 epochs with a simple model architecture.

Congratulations! You’ve built, trained, and assessed a simple CNN using PyTorch. Happy coding!

---

Here’s a special thanks to our [biggest supporters](https://www.aibutsimple.com/upgrade?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=convolutional-neural-networks-cnns-in-pytorch):

Sushant Waidande

---

If you enjoy our content, consider supporting us so we can keep doing what we do. Please share this with a friend!

Feedback, inquiries, advertising? Send us an email at [edwin@aibutsimple.com](mailto:edwin@aibutsimple.com).

If you like, you can also donate to our team to push out better newsletters every week!

[Support us ☕](https://ko-fi.com/aibutsimple?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=convolutional-neural-networks-cnns-in-pytorch)

---

That’s it for this week’s issue of AI, but simple. See you next week!

—AI, but simple team