Pytorch | Tech
by Saumya Shah
Some handy utilities
- You can always look at your model using summary:
from torchsummary import summary
model = resnet18(3, 1000)
summary(model.cuda(), (3, 224, 224))
- Saving and Loading Models
#save
torch.save(model.state_dict(), PATH)
#load
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
This rest of the post will cover some of the main blocks that we will use in Pytorch. For most purposes, the community is very helpful and the documentation of any PyTorch function is crystal clear. Not finding an answer is rare.
We shall do this by building a ResNet from scratch.
1. Module, Sequential, ModuleList and ModuleDict
These are useful to make your code more compact and modular. We aim to remove the repetitive parts as follows:
- Module
The Module is the main building block, it defines the base class for all neural network and you MUST subclass it.
import torch.nn as nn
import torch.nn.functional as F
class MyCNNClassifier(nn.Module):
def __init__(self, in_c, n_classes):
super().__init__()
self.conv1 = nn.Conv2d(in_c, 32, kernel_size=3, stride=1, padding=1)
self.bn1 = nn.BatchNorm2d(32)
self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
self.bn2 = nn.BatchNorm2d(32)
self.fc1 = nn.Linear(32 * 28 * 28, 1024)
self.fc2 = nn.Linear(1024, n_classes)
def forward(self, x):
x = self.conv1(x)
x = self.bn1(x)
x = F.relu(x)
x = self.conv2(x)
x = self.bn2(x)
x = F.relu(x)
x = x.view(x.size(0), -1) # flat
x = self.fc1(x)
x = F.sigmoid(x)
x = self.fc2(x)
return x
model = MyCNNClassifier(1, 10)
print(model)
# MyCNNClassifier( (conv1): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) (bn1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (conv2): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) (bn2): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (fc1): Linear(in_features=25088, out_features=1024, bias=True) (fc2): Linear(in_features=1024, out_features=10, bias=True) )
This is a basic classifier. If you wish to add another layer, you would have to add code in both init and forward function. To make it easier, we have sequential.
- Sequential
Sequential is a container of Modules that can be stacked together and run at the same time. While calling forward, the outputs are passed to next layer in sequential as we did in forward in previous case.
class MyCNNClassifier(nn.Module):
def __init__(self, in_c, n_classes):
super().__init__()
self.conv_block1 = nn.Sequential(
nn.Conv2d(in_c, 32, kernel_size=3, stride=1, padding=1),
nn.BatchNorm2d(32),
nn.ReLU()
)
self.conv_block2 = nn.Sequential(
nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1),
nn.BatchNorm2d(64),
nn.ReLU()
)
self.decoder = nn.Sequential(
nn.Linear(32 * 28 * 28, 1024),
nn.Sigmoid(),
nn.Linear(1024, n_classes)
)
def forward(self, x):
x = self.conv_block1(x)
x = self.conv_block2(x)
x = x.view(x.size(0), -1) # flat
x = self.decoder(x)
return x
model = MyCNNClassifier(1, 10)
print(model)
# MyCNNClassifier( (conv_block1): Sequential( (0): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (2): ReLU() ) (conv_block2): Sequential( (0): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (2): ReLU() ) (decoder): Sequential( (0): Linear(in_features=25088, out_features=1024, bias=True) (1): Sigmoid() (2): Linear(in_features=1024, out_features=10, bias=True) ) )
This is good, but we can go 2 steps further to reduce repetition of conv blocks by making a function and passing a list to sequential as follows:
def conv_block(in_f, out_f, *args, **kwargs):
return nn.Sequential(
nn.Conv2d(in_f, out_f, *args, **kwargs),
nn.BatchNorm2d(out_f),
nn.ReLU()
)
class MyCNNClassifier(nn.Module):
def __init__(self, in_c, n_classes):
super().__init__()
self.enc_sizes = [in_c, 32, 64]
conv_blocks = [conv_block(in_f, out_f, kernel_size=3, padding=1)
for in_f, out_f in zip(self.enc_sizes, self.enc_sizes[1:])]
self.encoder = nn.Sequential(*conv_blocks)
self.decoder = nn.Sequential(
nn.Linear(32 * 28 * 28, 1024),
nn.Sigmoid(),
nn.Linear(1024, n_classes)
)
def forward(self, x):
x = self.encoder(x)
x = x.view(x.size(0), -1) # flat
x = self.decoder(x)
return x
model = MyCNNClassifier(1, 10)
print(model)
# MyCNNClassifier( (encoder): Sequential( (0): Sequential( (0): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (2): ReLU() ) (1): Sequential( (0): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (2): ReLU() ) ) (decoder): Sequential( (0): Linear(in_features=25088, out_features=1024, bias=True) (1): Sigmoid() (2): Linear(in_features=1024, out_features=10, bias=True) ) )
NOTE: Sequential doesn’t accept list, we unpack it using the * operator.
The model can even be broken down into classes:
class MyEncoder(nn.Module):
def __init__(self, enc_sizes):
super().__init__()
self.conv_blokcs = nn.Sequential(*[conv_block(in_f, out_f, kernel_size=3, padding=1)
for in_f, out_f in zip(enc_sizes, enc_sizes[1:])])
def forward(self, x):
return self.conv_blokcs(x)
class MyDecoder(nn.Module):
def __init__(self, dec_sizes, n_classes):
super().__init__()
self.dec_blocks = nn.Sequential(*[dec_block(in_f, out_f)
for in_f, out_f in zip(dec_sizes, dec_sizes[1:])])
self.last = nn.Linear(dec_sizes[-1], n_classes)
def forward(self, x):
return self.dec_blocks()
class MyCNNClassifier(nn.Module):
def __init__(self, in_c, enc_sizes, dec_sizes, n_classes):
super().__init__()
self.enc_sizes = [in_c, *enc_sizes]
self.dec_sizes = [32 * 28 * 28, *dec_sizes]
self.encoder = MyEncoder(self.enc_sizes)
self.decoder = MyDecoder(dec_sizes, n_classes)
def forward(self, x):
x = self.encoder(x)
x = x.flatten(1) # flat
x = self.decoder(x)
return x
- ModuleList
This is useful when you need to store intermediate results or use them or manipulate before passing to next layer. It gives more control than sequential.
class MyModule(nn.Module):
def __init__(self, sizes):
super().__init__()
self.layers = nn.ModuleList([nn.Linear(in_f, out_f) for in_f, out_f in zip(sizes, sizes[1:])])
self.trace = []
def forward(self,x):
for layer in self.layers:
x = layer(x)
self.trace.append(x)
return x
- ModuleDict
This is useful when you wish to make changes in layers and shortens the code length.
def conv_block(in_f, out_f, activation='relu', *args, **kwargs):
activations = nn.ModuleDict([
['lrelu', nn.LeakyReLU()],
['relu', nn.ReLU()]
])
return nn.Sequential(
nn.Conv2d(in_f, out_f, *args, **kwargs),
nn.BatchNorm2d(out_f),
activations[activation]
)
2. ResNet
The motivation to use ResNet is to make deep networks. The conventional CNN networks didn’t get better as they became deeper because of vanishing gradients, i.e. the gradients became smaller as they were back propagated, and thus, initial layers didn’t learn well. The residual skipped connections solves that issue. But more than all of it, I chose this because there are some things about PyTorch that I can learn while building this.
Finally, we have what we needed to start with. Let’s make our ResNet now!
The flow is basically:
- Have a basic conv block
- Build a residual block skeleton class (with blocks and shortcut overriden later)
- Make a Resnet residual block from step 2, overriding the shortcut
- Make ResNet block from step 3, now overriding the blocks
- Introducing a Bottle Neck to reduce the number of parameters.
- Making a layer of blocks
- Finally making an encoder and a decoder
- Basic Convolution block
To make an auto-padding conv layer.
from functools import partial
class Conv2dAuto(nn.Conv2d):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.padding = (self.kernel_size[0] // 2, self.kernel_size[1] // 2) # dynamic add padding based on the kernel_size
conv3x3 = partial(Conv2dAuto, kernel_size=3, bias=False)
conv = conv3x3(in_channels=32, out_channels=64)
print(conv)
del conv
# Conv2dAuto(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
- Dictionary of activations
def activation_func(activation):
return nn.ModuleDict([
['relu', nn.ReLU(inplace=True)],
['leaky_relu', nn.LeakyReLU(negative_slope=0.01, inplace=True)],
['selu', nn.SELU(inplace=True)],
['none', nn.Identity()]
])[activation]
- Residual Block
The residual block takes an input with in_channels, applies some blocks of convolutional layers to reduce it to out_channels and sum it up to the original input. If their sizes mismatch, then the input goes into an identity. (NOTE: Identity is just a place holder which shall be over-ridden later)
class ResidualBlock(nn.Module):
def __init__(self, in_channels, out_channels, activation='relu'):
super().__init__()
self.in_channels, self.out_channels, self.activation = in_channels, out_channels, activation
self.blocks = nn.Identity()
self.activate = activation_func(activation)
self.shortcut = nn.Identity()
def forward(self, x):
residual = x
if self.should_apply_shortcut: residual = self.shortcut(x)
x = self.blocks(x)
x += residual
x = self.activate(x)
return x
@property
def should_apply_shortcut(self):
return self.in_channels != self.out_channels
In ResNet, each block has an expansion parameter in order to increase the out_channels if needed. Also, the identity is defined as a Convolution followed by a BatchNorm layer, this is referred to as shortcut. Then, we can just extend ResidualBlock and defined the shortcut function.
class ResNetResidualBlock(ResidualBlock):
def __init__(self, in_channels, out_channels, expansion=1, downsampling=1, conv=conv3x3, *args, **kwargs):
super().__init__(in_channels, out_channels, *args, **kwargs)
self.expansion, self.downsampling, self.conv = expansion, downsampling, conv
self.shortcut = nn.Sequential(
nn.Conv2d(self.in_channels, self.expanded_channels, kernel_size=1,
stride=self.downsampling, bias=False),
nn.BatchNorm2d(self.expanded_channels)) if self.should_apply_shortcut else None
@property
def expanded_channels(self):
return self.out_channels * self.expansion
@property
def should_apply_shortcut(self):
return self.in_channels != self.expanded_channels
- ResNet block
# Stack conv and bn
def conv_bn(in_channels, out_channels, conv, *args, **kwargs):
return nn.Sequential(conv(in_channels, out_channels, *args, **kwargs), nn.BatchNorm2d(out_channels))
# ResNet Block
class ResNetBasicBlock(ResNetResidualBlock):
"""
Basic ResNet block composed by two layers of 3x3conv/batchnorm/activation
"""
expansion = 1
def __init__(self, in_channels, out_channels, *args, **kwargs):
super().__init__(in_channels, out_channels, *args, **kwargs)
self.blocks = nn.Sequential(
conv_bn(self.in_channels, self.out_channels, conv=self.conv, bias=False, stride=self.downsampling),
activation_func(self.activation),
conv_bn(self.out_channels, self.expanded_channels, conv=self.conv, bias=False),
)
- BottleNeck
To increase the network depth while keeping the parameters size as low as possible, the authors defined a BottleNeck block that “The three layers are 1x1, 3x3, and 1x1 convolutions, where the 1×1 layers are responsible for reducing and then increasing (restoring) dimensions, leaving the 3×3 layer a bottleneck with smaller input/output dimensions.”
class ResNetBottleNeckBlock(ResNetResidualBlock):
expansion = 4
def __init__(self, in_channels, out_channels, *args, **kwargs):
super().__init__(in_channels, out_channels, expansion=4, *args, **kwargs)
self.blocks = nn.Sequential(
conv_bn(self.in_channels, self.out_channels, self.conv, kernel_size=1),
activation_func(self.activation),
conv_bn(self.out_channels, self.out_channels, self.conv, kernel_size=3, stride=self.downsampling),
activation_func(self.activation),
conv_bn(self.out_channels, self.expanded_channels, self.conv, kernel_size=1),
)
- ResNet Layer
A layer is composed of stacked blocks. We stack n blocks with the first block having downsampling with stride of 2.
NOTE : This doesn’t subclass the previous block.
class ResNetLayer(nn.Module):
"""
A ResNet layer composed by `n` blocks stacked one after the other
"""
def __init__(self, in_channels, out_channels, block=ResNetBasicBlock, n=1, *args, **kwargs):
super().__init__()
# 'We perform downsampling directly by convolutional layers that have a stride of 2.'
downsampling = 2 if in_channels != out_channels else 1
self.blocks = nn.Sequential(
block(in_channels , out_channels, *args, **kwargs, downsampling=downsampling),
*[block(out_channels * block.expansion,
out_channels, downsampling=1, *args, **kwargs) for _ in range(n - 1)]
)
def forward(self, x):
x = self.blocks(x)
return x
- Encoder
So ResNet is an Encoder-Decoder network. The Encoder is composed of layers with increasing features.
class ResNetEncoder(nn.Module):
"""
ResNet encoder composed by layers with increasing features.
"""
def __init__(self, in_channels=3, blocks_sizes=[64, 128, 256, 512], deepths=[2,2,2,2],
activation='relu', block=ResNetBasicBlock, *args, **kwargs):
super().__init__()
self.blocks_sizes = blocks_sizes
self.gate = nn.Sequential(
nn.Conv2d(in_channels, self.blocks_sizes[0], kernel_size=7, stride=2, padding=3, bias=False),
nn.BatchNorm2d(self.blocks_sizes[0]),
activation_func(activation),
nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
)
self.in_out_block_sizes = list(zip(blocks_sizes, blocks_sizes[1:]))
self.blocks = nn.ModuleList([
ResNetLayer(blocks_sizes[0], blocks_sizes[0], n=deepths[0], activation=activation,
block=block,*args, **kwargs),
*[ResNetLayer(in_channels * block.expansion,
out_channels, n=n, activation=activation,
block=block, *args, **kwargs)
for (in_channels, out_channels), n in zip(self.in_out_block_sizes, deepths[1:])]
])
def forward(self, x):
x = self.gate(x)
for block in self.blocks:
x = block(x)
return x
- Decoder
The last piece that has fully connected networks mapping features to classes.
class ResnetDecoder(nn.Module):
"""
This class represents the tail of ResNet. It performs a global pooling and maps the output to the
correct class by using a fully connected layer.
"""
def __init__(self, in_features, n_classes):
super().__init__()
self.avg = nn.AdaptiveAvgPool2d((1, 1))
self.decoder = nn.Linear(in_features, n_classes)
def forward(self, x):
x = self.avg(x)
x = x.view(x.size(0), -1)
x = self.decoder(x)
return x
- ResNet
Finally, we have our ResNet modules to build a network as below:
class ResNet(nn.Module):
def __init__(self, in_channels, n_classes, *args, **kwargs):
super().__init__()
self.encoder = ResNetEncoder(in_channels, *args, **kwargs)
self.decoder = ResnetDecoder(self.encoder.blocks[-1].blocks[-1].expanded_channels, n_classes)
def forward(self, x):
x = self.encoder(x)
x = self.decoder(x)
return x
# EXAMPLE RESNET 18
def resnet18(in_channels, n_classes, block=ResNetBasicBlock, *args, **kwargs):
return ResNet(in_channels, n_classes, block=block, deepths=[2, 2, 2, 2], *args, **kwargs)
References:
- https://towardsdatascience.com/pytorch-how-and-when-to-use-module-sequential-modulelist-and-moduledict-7a54597b5f17
- https://towardsdatascience.com/residual-network-implementing-resnet-a7da63c7b278