Writing ResNet from Scratch in PyTorch
Residual Networks, commonly known as ResNets, are revolutionary deep learning architectures introduced by Kaiming He et al. in 2015. They addressed the vanishing gradient problem, which had been a significant hurdle in training very deep neural networks. By introducing skip connections or residual blocks, ResNets make it possible to train networks with hundreds or even thousands of layers. In this blog post, we will guide you through the process of writing a ResNet from scratch using PyTorch, a popular deep learning framework.
Table of Contents#
- Understanding the Basics of ResNet
- Prerequisites
- Setting up the Environment
- Building the Residual Block
- Constructing the ResNet Architecture
- Training the ResNet Model
- Evaluating the Model
- Conclusion
- References
1. Understanding the Basics of ResNet#
Vanishing Gradient Problem#
In a deep neural network, during backpropagation, gradients can become extremely small as they are passed through multiple layers. This makes it difficult for the network to learn effectively, especially for deeper layers.
Residual Blocks#
ResNets solve the vanishing gradient problem by introducing residual blocks. In a residual block, the input to the block is added to the output of the block's convolutional layers. This skip connection allows the network to learn the residual mapping, which is often easier to optimize than the original mapping.
The mathematical formula for a residual block is: [F(x) + x] where (F(x)) is the output of the convolutional layers in the block and (x) is the input to the block.
2. Prerequisites#
- Basic knowledge of Python programming.
- Understanding of neural networks, convolutional neural networks (CNNs), and backpropagation.
- Familiarity with PyTorch, including tensors, neural network modules, and training loops.
3. Setting up the Environment#
Installing PyTorch#
You can install PyTorch using pip or conda depending on your preference. For example, to install PyTorch with CPU support using pip, you can run the following command:
pip install torch torchvisionImporting Libraries#
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np4. Building the Residual Block#
class ResidualBlock(nn.Module):
def __init__(self, in_channels, out_channels, stride = 1, downsample = None):
super(ResidualBlock, self).__init__()
self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size = 3, stride = stride, padding = 1, bias = False)
self.bn1 = nn.BatchNorm2d(out_channels)
self.relu = nn.ReLU(inplace = True)
self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size = 3, stride = 1, padding = 1, bias = False)
self.bn2 = nn.BatchNorm2d(out_channels)
self.downsample = downsample
def forward(self, x):
residual = x
out = self.conv1(x)
out = self.bn1(out)
out = self.relu(out)
out = self.conv2(out)
out = self.bn2(out)
if self.downsample:
residual = self.downsample(x)
out += residual
out = self.relu(out)
return outExplanation#
__init__: Initializes the layers of the residual block. Thedownsampleargument is used when the input and output channels or the spatial dimensions of the block do not match.forward: Defines the forward pass of the block. The inputxis passed through two convolutional layers with batch normalization and ReLU activation. Ifdownsampleis notNone, the input is downsampled. Finally, the output of the convolutional layers is added to the residual (either the original input or the downsampled input) and passed through another ReLU activation.
5. Constructing the ResNet Architecture#
class ResNet(nn.Module):
def __init__(self, block, layers, num_classes = 10):
super(ResNet, self).__init__()
self.in_channels = 64
self.conv1 = nn.Conv2d(3, 64, kernel_size = 7, stride = 2, padding = 3, bias = False)
self.bn1 = nn.BatchNorm2d(64)
self.relu = nn.ReLU(inplace = True)
self.maxpool = nn.MaxPool2d(kernel_size = 3, stride = 2, padding = 1)
self.layer1 = self._make_layer(block, 64, layers[0])
self.layer2 = self._make_layer(block, 128, layers[1], stride = 2)
self.layer3 = self._make_layer(block, 256, layers[2], stride = 2)
self.layer4 = self._make_layer(block, 512, layers[3], stride = 2)
self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
self.fc = nn.Linear(512, num_classes)
def _make_layer(self, block, out_channels, blocks, stride = 1):
downsample = None
if stride != 1 or self.in_channels != out_channels:
downsample = nn.Sequential(
nn.Conv2d(self.in_channels, out_channels, kernel_size = 1, stride = stride, bias = False),
nn.BatchNorm2d(out_channels))
layers = []
layers.append(block(self.in_channels, out_channels, stride, downsample))
self.in_channels = out_channels
for i in range(1, blocks):
layers.append(block(out_channels, out_channels))
return nn.Sequential(*layers)
def forward(self, x):
out = self.conv1(x)
out = self.bn1(out)
out = self.relu(out)
out = self.maxpool(out)
out = self.layer1(out)
out = self.layer2(out)
out = self.layer3(out)
out = self.layer4(out)
out = self.avgpool(out)
out = out.view(out.size(0), -1)
out = self.fc(out)
return outExplanation#
__init__: Initializes the overall ResNet architecture. It includes an initial convolutional layer, batch normalization, ReLU activation, max pooling, four residual layers, average pooling, and a fully connected layer._make_layer: Creates a sequence of residual blocks. It checks if downsampling is required and adds the appropriate downsampling layer.forward: Defines the forward pass of the entire ResNet model. The input is passed through the initial layers, four residual layers, average pooling, and the fully connected layer.
6. Training the ResNet Model#
Loading the Dataset#
transform = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4,
shuffle=True, num_workers=2)
testset = torchvision.datasets.CIFAR10(root='./data', train=False,
download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=4,
shuffle=False, num_workers=2)
classes = ('plane', 'car', 'bird', 'cat',
'deer', 'dog', 'frog', 'horse', 'ship', 'truck')Initializing the Model and Optimizer#
model = ResNet(ResidualBlock, [2, 2, 2, 2])
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)Training Loop#
for epoch in range(2): # loop over the dataset multiple times
running_loss = 0.0
for i, data in enumerate(trainloader, 0):
# get the inputs; data is a list of [inputs, labels]
inputs, labels = data
# zero the parameter gradients
optimizer.zero_grad()
# forward + backward + optimize
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
# print statistics
running_loss += loss.item()
if i % 2000 == 1999: # print every 2000 mini-batches
print(f'[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 2000:.3f}')
running_loss = 0.0
print('Finished Training')7. Evaluating the Model#
correct = 0
total = 0
with torch.no_grad():
for data in testloader:
images, labels = data
outputs = model(images)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print(f'Accuracy of the network on the 10000 test images: {100 * correct / total} %')8. Conclusion#
In this blog post, we have walked through the process of writing a ResNet from scratch in PyTorch. We started by understanding the basic concept of ResNets, then built the residual block and the entire ResNet architecture. Finally, we trained and evaluated the model on the CIFAR - 10 dataset. ResNets are a powerful tool in deep learning, and implementing them from scratch can help you gain a deeper understanding of their inner workings.
9. References#
- He, K., Zhang, X., Ren, S., & Sun, J. (2016). Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition (pp. 770-778).
- PyTorch official documentation: https://pytorch.org/docs/stable/index.html
- CIFAR - 10 dataset: https://www.cs.toronto.edu/~kriz/cifar.html