Creating a Neural Network in Pytorch

One of the most fundamental functionalities of any Deep Learning library is to create a Neural Network. In this chapter of the Pytorch Tutorial, you will learn how to create a Neural Network using Pytorch. Specifically, you will learn how to create a dense feed-forward Artificial Neural Network(ANN) model in Pytorch. However, you can also create other types of Neural Networks in a similar fashion.

To create a Neural Network, you must create a class for that Neural Network and then instantiate that class.

Importing Modules

In the examples below, we have imported have made use of modules and sub-modules from the Pytorch library. Therefore, it is important to import them.

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

Subclassing the nn.Module

To create your own model in Pytorch, you have to subclass the nn.Module class. You have to implement the constructor and the forward() method in your model class. The architecture of the model is defined inside the constructor of the class. Whereas, the flow of data through the model is defined inside the forward() method of the class.

Below is the skeleton code for a Neural Network class.

class MyNeuralNetwork(nn.Module):
    def __init__(self):
        #code for the architecture of the Neural Network
        pass
    def forward(self, x):
        #code for the flow of data through the Neural Network
        pass

Architecture of Model

The architecture of the model is defined in the constructor of the Neural Network class. The layers of a Neural Networks are added as attributes inside the constructor.

Example

In this example, we will build the architecture of MyNeuralNetwork class.

Firstly, we are using the constructor of the parent class, nn.Module to initialize the parameters. We pass the class, MyNeuralNetwork along with self– which will refer to the object(the Neural Network) which will be created using this class to the constructor of nn.Module. This is always the first step to do inside the constructor of a Neural Network class.

Additionally, we have also created the architecture of the Neural Network which consists of 3 dense layers-

dense_layer1 – The first dense layer having 784 inputs and 300 outputs.
dense_layer2 – The second dense layer having 300 inputs and 100 outputs.
dense_layer3 – The third dense layer having 100 inputs and 10 outputs.

def __init__(self):
    super(MyNeuralNetwork, self).__init__()
    self.dense_layer1 = nn.Linear(784, 300)
    self.dense_layer2 = nn.Linear(300, 100)
    self.dense_layer3 = nn.Linear(100, 10)

Data Flow through the Model

The flow of data through the model is defined in the forward() method of the Neural Network class.

Example

In this example, we will define the flow of data through the architecture of the MyNeuralNetwork class that we defined in the earlier example.

We are applying the ReLU activation function to the outputs of dense_layer1 and dense_layer2, and the Softmax activation function to the output of dense_layer3. Optionally, we can also perform operations such as reshaping the input tensor to the desired shape, using loops, etc inside the forward() method.

def forward(self, x):

    # applying ReLU activation function to the output of dense_layer1
    x = F.relu(self.dense_layer1)

    # applying ReLU activation function to the output of dense_layer2
    x = F.relu(self.dense_layer2)

    # applying softmax activation function to the output of dense_layer3
    x = F.softmax(self.dense_layer3)

    #returning the output
    return x

Creating the Model

Once you have created the class for your Neural Network, you just need to instantiate that class to create your model. This will create a Neural Network.

Example

In this example, we are creating a Neural Network model, mynet by instantiating the MyNeuralNetwork class.

# Creating a Neural Network model
mynet = MyNeuralNetwork()