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
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.
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).
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
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
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
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%
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.
Table of Content
- Why Custom Layers?
- Building The Custom Layer
- Creating a Custom Network
- The Main Program
- Conclusion