Skip to content
Related Articles

Related Articles

How to Define a Simple Convolutional Neural Network in PyTorch?

Improve Article
Save Article
  • Last Updated : 23 Sep, 2022
Improve Article
Save Article

In this article, we are going to see how to  Define a Simple Convolutional Neural Network in PyTorch using Python. Convolutional Neural Networks(CNN) is a type of Deep Learning algorithm which is highly instrumental in learning patterns and features in images. CNN has a unique trait which is its ability to process data with a grid-like topology whereas a typical Artificial Neural Network(Dense or Sparse) generally takes input by flattening the tensors into a one-dimensional vector. This facilitates it to learn and differentiate between features in images, which when represented digitally are essentially a grid of numbers.

Convolutional Neural Networks are typically comprised of multiple layers. Usually, the initial layers are used to detect simple features such as edges, and complex features are detected down the line, as we go deeper into the network.

CNN has countless qualities that make it so suitable for processing images. Let’s take a look at some of them:-

  • They require much less data pre-processing than other Deep Learning Algorithms.
  • A  well-trained CNN model has the ability to learn and classify features in an image, which gives much better accuracy in the classification and detection of features in images.
  • It can save a lot of computational resources by methods like increasing the convolutional and pooling layers.

What are Convolutional and Pooling Layers in CNN?

Convolutional Layers:

These are the first layers in a CNN, and they can be thought of as “Filters” for an image. Just like Filters in Instagram detect our face, a convolutional layer detects features or filters such as edges in an image, wherever they might be present.

Pooling Layers:

The pooling layers mainly reduce the computational cost by reducing the spatial size of the image. The best way to describe it would be that it makes the grids of information smaller by taking a “lump-sum” of the images’ spatial resolution.

Defining a Convolutional Neural Network using PyTorch:

Let’s now move on to define a simple Convolutional Neural Network with one Convolutional Layer and one Linear  Layer.

Step 1: Import the necessary libraries to define our own Convolutional Neural Network.

We will import the Torch library first. If not installed, just use the following pip command for the same:-

pip install torch

For this CNN model, we will be using the Adam optimizer, and ReLU inplace activation function.


# Step 1 is importing all the
# necessary libraries needed for defining our CNN
import torch
from torch.autograd import Variable
from torch.nn import (Linear, ReLU,
                      Sequential, Conv2d,
                      MaxPool2d, Module,
                      Softmax, BatchNorm2d,
from torch.optim import Adam, SGD

Step 2: We will start by defining the class Net() to build a constructor for our CNN of the desired shape.

Notice in Step 1, We have imported a module torch.nn. This module contains different classes that can help in building neural networks.


# Step 2 is defining our CNN
# it's components including the activation functions,
# CNN layers then it's
class Net(Module):
    def __init__(self):
        super(Net, self).__init__()
        # Now we will define a simple CNN
        # with only one 2D Convolutional Layer
        # The layer has it's own inplace
        # Relu activation function
        self.cnn_layers = Sequential(
            # Declaring the Convolutional Layer
            Conv2d(1, 4, kernel_size=3, stride=1, padding=1),
            MaxPool2d(kernel_size=2, stride=2),
        self.linear_layers = Sequential(
            Linear(4 * 8 * 8, 10)
    # Defining the Forward feed into the linear
    # layers from the Convolutional Layers.
    def forward(self, x):
        x = self.cnn_layers(x)
        x = x.view(x.size(0), -1)
        x = self.linear_layers(x)
        return x

Step 3: The last step will be defining a model of our choice and printing it’s attributes.

We can easily define any number of models after creating the above-shown class by just calling the constructor function and we will obtain a Convolutional Neural Network with two 2D Convolutional Layers and one Linear Layer.

The print(model) line gives the shape of the CNN as it’s output.


# Defining the CNN model using the 
# constructed class
model = Net()
# Adam is used as an optimizer for
# the model parameters
optimizer = Adam(model.parameters(), lr=0.07)
# CrossEntropyLoss is chosen as the
# criteria to evaluate the model
# (loss function)
criterion = CrossEntropyLoss()
# CUDA tensor types use GPU instead of
# CPUs for computation the following 
# function checks whether the GPU is available 
# If the GPU is available,the model
# uses GPU for computation
if torch.cuda.is_available():
    model = model.cuda()
    criterion = criterion.cuda()

This code will define a simple CNN in PyTorch which uses a single Convolutional Layer and a single Linear Layer.



  (cnn_layers): Sequential(

    (0): Conv2d(1, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))

    (1): BatchNorm2d(4, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)

    (2): ReLU(inplace=True)

    (3): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)


  (linear_layers): Sequential(

    (0): Linear(in_features=256, out_features=10, bias=True)



Applications of a CNN classifier:

  • CNN-based classifiers can be used for object and pattern detection for countless purposes, e.g. face recognition, Classification, and Regression prediction problems.
  • It’s really beneficial in Cancer Detection and Biometric Authentication. CNNs have had a huge impact on healthcare and disease detection.
  • Well trained CNNs can be used to caption images.

My Personal Notes arrow_drop_up
Related Articles

Start Your Coding Journey Now!