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.
- __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.
- 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
