Create Custom Neural Network in PyTorch

PyTorch is a popular deep learning framework, empowers you to build and train powerful neural networks. But what if you need to go beyond the standard layers offered by the library? Here's where custom layers come in, allowing you to tailor the network architecture to your specific needs. This comprehensive guide explores how to create custom layers in PyTorch, unlocking a new level of flexibility for your deep learning projects.

Why Custom Layers?

While PyTorch provides a rich set of built-in layers, there are scenarios where you might require more specialized functionality. Custom layers enable you to:

• Implement unique operations: Design layers for specific tasks not covered by existing layers.
• Incorporate domain knowledge: Integrate domain-specific knowledge into your neural network architecture.
• Experiment with novel architectures: Prototype and test innovative network structures for research purposes.

By creating custom layers, you gain greater control over your deep learning models, pushing the boundaries of what's achievable.

Building The Custom Layer

Let's dive into the practical aspects of creating a custom layer in PyTorch. We'll start with a simple example that performs element-wise multiplication.

• Inheritance: The CustomLayer class inherits from nn.Module, the foundation for building neural network layers in PyTorch.
• Initialization: The __init__ method defines the layer's parameters. Here, we create a weight tensor with the same size as the input (ip_size).
• Weight Parameter: The weight tensor is converted to a learnable parameter using nn.Parameter allowing the network to optimize the weights during training.
• Weight Initialization: We initialize the weights using a normal distribution with a mean of 0 and a standard deviation of 1.
• Forward Pass: The forward method defines the core operation of the layer. In this case, it performs element-wise multiplication between the input x and the weight parameter self.weight.

Python3

import torch
import torch.nn as nn

class CustomLayer(nn.Module):
  def __init__(self,ip_size):
    super().__init__()
    self.size = ip_size
    weight_tensor = torch.Tensor(self.size)
    self.weight = nn.Parameter(weight_tensor)
    torch.nn.init.normal_(self.weight,mean=0.0,std=1.0)

  def forward(self,x):
    return x * self.weight

Creating a Custom Network

Now that you have a custom layer, let's see how to use it within a neural network. An example that stacks your CustomLayer with a ReLU activation.

This code defines a simple network with two layers:

1. The first layer is your custom CustomLayer with an input size of 10.
2. The second layer is a ReLU activation layer from the nn module.

Python3

CustomNetwork = nn.Sequential(
    CustomLayer(10),
    nn.ReLU()
)

The Main Program

Let's put together all steps and see an example that demonstrates building a custom linear layer and training an Artificial Neural Network (ANN) for image classification.

1. Import libraries

Python3

import torch
import torch.nn as nn
import torch.nn.functional as F
import math
import numpy

2) Define your custom layer class

This code defines a CustomLinear layer that mimics the behavior of a fully connected layer in PyTorch.

CustomLinear offers a custom implementation of a linear layer with learnable weights and biases, performing linear transformations during the network's forward pass.

Python3

class CustomLinear(nn.Module):
    def __init__(self, in_features, out_features):
        super(CustomLinear, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.weight = nn.Parameter(torch.Tensor(out_features, in_features))
        self.bias = nn.Parameter(torch.Tensor(out_features))
        self.reset_parameters()

    def reset_parameters(self):
        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
        if self.bias is not None:
            fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight)
            bound = 1 / math.sqrt(fan_in)
            nn.init.uniform_(self.bias, -bound, bound)

    def forward(self, input):
        return input.matmul(self.weight.t()) + self.bias

3) Use your custom layer in a neural network

CustomNetwork is a neural network composed of multiple custom linear layers (CustomLinear). The __init__ method initializes the network with five custom linear layers. The forward method defines the forward pass of the network:

• It reshapes the input tensor x into a 2D tensor (x.view(-1, 28*28)).
• It applies the custom linear layers (fc1, fc2, fc3, fc4, fc5) sequentially, followed by ReLU activation functions (F.relu).

Python3

class CustomNetwork(nn.Module):
    def __init__(self):
        super(CustomNetwork, self).__init__()
        self.fc1 = CustomLinear(28*28, 512)
        self.fc2 = CustomLinear(512, 256)
        self.fc3 = CustomLinear(256, 128)
        self.fc4 = CustomLinear(128, 64)
        self.fc5 = CustomLinear(64, 10)

    def forward(self, x):
        x = x.view(-1, 28*28)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = F.relu(self.fc3(x))
        x = F.relu(self.fc4(x))
        x = self.fc5(x)
        return x

4) Instantiate your network and test it

Python3

model = CustomNetwork()

from torchsummary import summary
summary(model, (1, 28, 28))
dummy_input = torch.randn(1, 1, 28, 28)

Output:

----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
      CustomLinear-1                  [-1, 512]         401,920
      CustomLinear-2                  [-1, 256]         131,328
      CustomLinear-3                  [-1, 128]          32,896
      CustomLinear-4                   [-1, 64]           8,256
      CustomLinear-5                   [-1, 10]             650
================================================================
Total params: 575,050
Trainable params: 575,050
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.00
Forward/backward pass size (MB): 0.01
Params size (MB): 2.19
Estimated Total Size (MB): 2.20
----------------------------------------------------------------

5) Dataset load - Training and optimization

Python3

from torchvision import datasets, transforms
transform = transforms.Compose([transforms.ToTensor(),
                                transforms.Normalize((0.5,), (0.5,))])

trainset = datasets.MNIST('~/.pytorch/MNIST_data/', download=True, train=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True)

testset = datasets.MNIST('~/.pytorch/MNIST_data/', download=True, train=False, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=64, shuffle=False)

import torch.optim as optim

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

epochs = 5
for e in range(epochs):
    running_loss = 0
    for images, labels in trainloader:
        optimizer.zero_grad()
        output = model(images)
        loss = criterion(output, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    else:
        print(f"Training loss: {running_loss/len(trainloader)}")

Output:

Training loss: 0.3632937747079617
Training loss: 0.15736598590972709
Training loss: 0.11662620213974927
Training loss: 0.09595851497679973
Training loss: 0.08306442246748023

6.) Accuracy

Python3

correct = 0
total = 0
with torch.no_grad():
    for images, labels in testloader:
        output = model(images)
        _, predicted = torch.max(output.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f"Accuracy: {100 * correct / total}%")

Output:

Accuracy: 97.29%

Conclusion

In conclusion, creating a custom layer in PyTorch allows you to extend the functionality of the framework and implement custom computations tailored to your specific needs.