Create Model using Custom Module in Pytorch

Custom module in Pytorch

A custom module in PyTorch is a user-defined module that is built using the PyTorch library’s built-in neural network module, torch.nn.Module. It’s a way of creating new modules by combining and extending the functionality provided by existing PyTorch modules.

The torch.nn.Module class provides a convenient way to create custom modules because it includes some key features that are important for building neural networks, such as the ability to keep track of learnable parameters and the ability to perform automatic differentiation (for computing gradients during training).

By creating a new class that inherits from torch.nn.Module, and defining an __init__ method to initialize the module’s parameters, and forward method that perform the computation, we can create our own custom module. These custom modules can be used just like any of the built-in PyTorch modules, such as torch.nn.Module or torch.nn.Conv2d, and can be included in a larger model architecture.

Creating a custom module can be useful in many situations. For example, we might create a custom module to implement a novel layer or activation function that is not included in PyTorch’s built-in modules. Or we could create a custom module that represents a more complex model, such as a sequence-to-sequence model, composed of multiple layers and other modules.

When creating a custom data model using a custom module in PyTorch, we will need to define a subclass of the torch.nn.Module class and define the __init__() and forward() methods.

1. __init__(): The __init__ method is used to initialize the module’s parameters. This method is called when the module is created, and it allows we to set up any internal state that the module needs. For example, we might use this method to initialize the weights of a neural network or to create other modules that the module needs in order to function.
2. forward(): The forward method is used to perform the computation that the module represents. This method takes in one or more input tensors, performs computations on them, and returns the output tensors. It is a forward pass of the module.

Once defined the custom module, we can create an instance of the module and use it to train a model by defining the loss function and optimizer, and then iterating through the training data to perform the forward and backward passes and optimize the model parameters.

Before going forward with creating a custom module in Pytorch, we have to install the torch library using the following command:

pip install torch

Here is a step-by-step example of creating a custom module in PyTorch and training it on a dataset from torchvision.datasets:

Step 1: Import the necessary libraries

Python3

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from sklearn.metrics import classification_report

Step 2: Create the module

In this step, we define a custom module called MyModule by creating a new class that inherits from the nn.Module base class. In the __init__ method, define the architecture of the model by creating the necessary layers. Here, we create two linear layers, one with num_inputs and hidden_size, and the other one with hidden_size and num_outputs.

Defining the forward pass: Here we define the forward pass of the model within the class by implementing the forward method. In this example, the input is passed through the first linear layer, then a relu activation function is applied to it, and then it is passed through the second linear layer.

Python3

class MyModule(nn.Module):

    # Initialize the parameter
    def __init__(self, num_inputs, num_outputs, hidden_size):
        super(MyModule, self).__init__()
        self.linear1 = nn.Linear(num_inputs, hidden_size)
        self.linear2 = nn.Linear(hidden_size, num_outputs)

    # Forward pass
    def forward(self, input):
        lin    = self.linear1(input)
        output = nn.functional.relu(lin)
        pred   = self.linear2(output)
        return pred

Step 3: Define the loss function and optimizer

Python3

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(my_module.parameters(), lr=0.005)

 Step 4: Define the transformations for the dataset

Python3

transform = transforms.Compose([transforms.ToTensor(),
                                transforms.Normalize((0.5,), (0.5,))])

Step 5: Load the dataset

In this step, we load the MNIST dataset with pytorch dataset from torchvision.datasets. This will download the datasets and save them in the data folder. Here we use transform to transform the dataset into pytorch tensor.

Python3

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

Step 5: Define the Dataloader

This will convert the dataset into batch size of 64.

Python3

# Define the data loader
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=64, shuffle=False)

Step 6: Training the model

Python3

for epoch in range(20):
    for i, (images, labels) in enumerate(train_loader):
        images = images.view(-1, 28*28)
        optimizer.zero_grad()
        output = my_module(images)
        loss = criterion(output, labels)
        loss.backward()
        optimizer.step()
    print('Epoch',epoch,' Loss -->',loss)

Output:

Epoch 0  Loss --> tensor(0.7018, grad_fn=<NllLossBackward0>)
Epoch 1  Loss --> tensor(0.7541, grad_fn=<NllLossBackward0>)
Epoch 2  Loss --> tensor(0.2634, grad_fn=<NllLossBackward0>)
Epoch 3  Loss --> tensor(0.7319, grad_fn=<NllLossBackward0>)
Epoch 4  Loss --> tensor(0.4635, grad_fn=<NllLossBackward0>)
Epoch 5  Loss --> tensor(0.1691, grad_fn=<NllLossBackward0>)
Epoch 6  Loss --> tensor(0.3424, grad_fn=<NllLossBackward0>)
Epoch 7  Loss --> tensor(0.3373, grad_fn=<NllLossBackward0>)
Epoch 8  Loss --> tensor(0.4458, grad_fn=<NllLossBackward0>)
Epoch 9  Loss --> tensor(0.3922, grad_fn=<NllLossBackward0>)
Epoch 10  Loss --> tensor(0.3963, grad_fn=<NllLossBackward0>)
Epoch 11  Loss --> tensor(0.4633, grad_fn=<NllLossBackward0>)
Epoch 12  Loss --> tensor(0.2869, grad_fn=<NllLossBackward0>)
Epoch 13  Loss --> tensor(0.3894, grad_fn=<NllLossBackward0>)
Epoch 14  Loss --> tensor(0.0586, grad_fn=<NllLossBackward0>)
Epoch 15  Loss --> tensor(0.5466, grad_fn=<NllLossBackward0>)
Epoch 16  Loss --> tensor(0.2478, grad_fn=<NllLossBackward0>)
Epoch 17  Loss --> tensor(0.2898, grad_fn=<NllLossBackward0>)
Epoch 18  Loss --> tensor(0.0774, grad_fn=<NllLossBackward0>)
Epoch 19  Loss --> tensor(0.2922, grad_fn=<NllLossBackward0>)

Step 7: Model Evaluation

Once we have trained our custom model on a dataset, we can use it to make predictions and evaluate its performance using a classification report. A classification report is a summary of the performance of a classification model and includes several metrics such as precision, recall, f1-score, and support.

Python3

# Test the model
with torch.no_grad():
    y_true = []
    y_pred = []
    correct = 0
    total = 0
    for images, labels in test_loader:
        images = images.view(-1, 28*28)
        output = my_module(images)
        _, predicted = torch.max(output.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum()
        y_true += labels.tolist()
        y_pred += predicted.tolist()

    # Accuracy
    print('Accuracy: {} %'.format(100 * correct / total))

    # Generate the classification report
    report = classification_report(y_true, y_pred)
    print(report)

Output:

Accuracy: 92.94999694824219 %
              precision    recall  f1-score   support

           0       0.95      0.98      0.96       980
           1       0.96      0.97      0.97      1135
           2       0.92      0.92      0.92      1032
           3       0.93      0.90      0.91      1010
           4       0.93      0.93      0.93       982
           5       0.88      0.90      0.89       892
           6       0.93      0.94      0.94       958
           7       0.95      0.92      0.94      1028
           8       0.91      0.90      0.90       974
           9       0.91      0.92      0.92      1009

    accuracy                           0.93     10000
   macro avg       0.93      0.93      0.93     10000
weighted avg       0.93      0.93      0.93     10000

Complete code

Python3

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from sklearn.metrics import classification_report

class MyModule(nn.Module):
    def __init__(self, num_inputs, num_outputs, hidden_size):
        super(MyModule, self).__init__()
        self.linear1 = nn.Linear(num_inputs, hidden_size)
        self.linear2 = nn.Linear(hidden_size, num_outputs)

    def forward(self, input):
        lin    = self.linear1(input)
        output = nn.functional.relu(lin)
        pred   = self.linear2(output)
        return pred

# Instantiate the custom module
my_module = MyModule(num_inputs=28*28, num_outputs=10, hidden_size=20)

# Define the loss function and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(my_module.parameters(), lr=0.01)

# Define the transformations for the dataset
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])

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

# Define the data loader
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=64, shuffle=False)

# Train the model
for epoch in range(10):
    for i, (images, labels) in enumerate(train_loader):
        images = images.view(-1, 28*28)
        optimizer.zero_grad()
        output = my_module(images)
        loss = criterion(output, labels)
        loss.backward()
        optimizer.step()
    print('Epoch -->',epoch,'-->',loss)



#Test the model
with torch.no_grad():
    y_true = []
    y_pred = []
    correct = 0
    total = 0
    for images, labels in test_loader:
        images = images.view(-1, 28*28)
        output = my_module(images)
        _, predicted = torch.max(output.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum()
        y_true += labels.tolist()
        y_pred += predicted.tolist()

    # Accuracy
    print('Accuracy: {} %'.format(100 * correct / total))

    # Classification Report
    report = classification_report(y_true, y_pred)
    print(report)

Output

Epoch 0  Loss --> tensor(0.3960, grad_fn=<NllLossBackward0>)
Epoch 1  Loss --> tensor(0.3620, grad_fn=<NllLossBackward0>)
Epoch 2  Loss --> tensor(0.5146, grad_fn=<NllLossBackward0>)
Epoch 3  Loss --> tensor(0.6453, grad_fn=<NllLossBackward0>)
Epoch 4  Loss --> tensor(0.1580, grad_fn=<NllLossBackward0>)
Epoch 5  Loss --> tensor(0.3488, grad_fn=<NllLossBackward0>)
Epoch 6  Loss --> tensor(0.2321, grad_fn=<NllLossBackward0>)
Epoch 7  Loss --> tensor(0.1614, grad_fn=<NllLossBackward0>)
Epoch 8  Loss --> tensor(0.1180, grad_fn=<NllLossBackward0>)
Epoch 9  Loss --> tensor(0.6431, grad_fn=<NllLossBackward0>)
Accuracy: 93.52999877929688 %
              precision    recall  f1-score   support

           0       0.95      0.97      0.96       980
           1       0.96      0.98      0.97      1135
           2       0.93      0.92      0.93      1032
           3       0.92      0.92      0.92      1010
           4       0.94      0.93      0.93       982
           5       0.92      0.90      0.91       892
           6       0.92      0.96      0.94       958
           7       0.94      0.93      0.94      1028
           8       0.92      0.91      0.91       974
           9       0.93      0.92      0.93      1009

    accuracy                           0.94     10000
   macro avg       0.93      0.93      0.93     10000
weighted avg       0.94      0.94      0.94     10000