import torch
import torch.nn as nn
from torchinfo import summary
import torchvision.models17 UNet Encoders
17.1 Introduction
The backbone or encoder component of semantic segmentation models serve the same purpose as the convolutional layers in a CNN architecture designed for scene labeling tasks: they characterize spatial patterns at multiple spatial scales. As a result, the encoder component of semantic segmentation architectures, such as UNets, can be augmented to use a variety of pre-defined CNN architectures, such as ResNets and VGGNets. Since these common architectures have been trained using large datasets, such as ImageNet, this allows for pre-trained weights to also be incorporated into the backbone or encoder component of semantic segmentation models. This component of the network can then either be frozen or updated during the training process. Such a use of transfer learning may allow for training models with less training data and/or for few epochs to obtain adequate results.
In this chapter, I will demonstrate augmenting UNet. The first example will use a VGGNet-16 architecture as the model backbone, the second example will use a ResNet backbone, and the third example will use the efficient MobileNetV2 architecture. I will only define the model architectures and summarize them using torchinfo. I will not train the algorithms. However, the methods used in the prior chapter could be applied to these architectures if desired.
Since I am defining architectures, I need to import torch and torch.nn. The torchinfo package is used to summarize the model architecture while the backbones will be accessed using the torchvision implementation. Lastly, I define the GPU as the device.
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)cuda:017.2 VGGNet-16 Encoder
A UNet with a VGGNet-16 encoder is conceptualized in Figure 17.1. The convolutional blocks of VGGNet-16 serve as the encoder. The output of the last convolutional layer in each of the first four blocks, before max pooling is applied, is passed to the decoder through a skip connection, and the output of the fifth block is passed through max pooling and a bottleneck. The decoder is the same as the one used in Chapter 15: each block applies 2D transpose convolution, concatenates the upsampled feature maps with the encoder feature maps, and applies two 3x3 convolutional layers.
I begin by creating a VGGNet backbone. The first step is to instantiate the backbone, as implemented in the torchvision.models subpackage. I use a modified version of VGGNet-16 that incorporates batch normalization using the vgg16_bn() function. Since I will not actually train the algorithm, I do not download the pre-trained weights.
vgg16 = torchvision.models.vgg16_bn(weights=None).to(device)I next obtain a summary of the VGGNet-16 architecture using an input array size of (12, 3,256,256). In order to use this architecture within UNet, I will need to be able to obtain the outputs before each max pooling operation so that the results can be passed to the decoder via the skip connections. The array sizes in the spatial dimensions being passed through the skip connections must be the same as those to which they are being concatenated.
summary(vgg16, (12,3,256,256))==========================================================================================
Layer (type:depth-idx) Output Shape Param #
==========================================================================================
VGG [12, 1000] --
├─Sequential: 1-1 [12, 512, 8, 8] --
│ └─Conv2d: 2-1 [12, 64, 256, 256] 1,792
│ └─BatchNorm2d: 2-2 [12, 64, 256, 256] 128
│ └─ReLU: 2-3 [12, 64, 256, 256] --
│ └─Conv2d: 2-4 [12, 64, 256, 256] 36,928
│ └─BatchNorm2d: 2-5 [12, 64, 256, 256] 128
│ └─ReLU: 2-6 [12, 64, 256, 256] --
│ └─MaxPool2d: 2-7 [12, 64, 128, 128] --
│ └─Conv2d: 2-8 [12, 128, 128, 128] 73,856
│ └─BatchNorm2d: 2-9 [12, 128, 128, 128] 256
│ └─ReLU: 2-10 [12, 128, 128, 128] --
│ └─Conv2d: 2-11 [12, 128, 128, 128] 147,584
│ └─BatchNorm2d: 2-12 [12, 128, 128, 128] 256
│ └─ReLU: 2-13 [12, 128, 128, 128] --
│ └─MaxPool2d: 2-14 [12, 128, 64, 64] --
│ └─Conv2d: 2-15 [12, 256, 64, 64] 295,168
│ └─BatchNorm2d: 2-16 [12, 256, 64, 64] 512
│ └─ReLU: 2-17 [12, 256, 64, 64] --
│ └─Conv2d: 2-18 [12, 256, 64, 64] 590,080
│ └─BatchNorm2d: 2-19 [12, 256, 64, 64] 512
│ └─ReLU: 2-20 [12, 256, 64, 64] --
│ └─Conv2d: 2-21 [12, 256, 64, 64] 590,080
│ └─BatchNorm2d: 2-22 [12, 256, 64, 64] 512
│ └─ReLU: 2-23 [12, 256, 64, 64] --
│ └─MaxPool2d: 2-24 [12, 256, 32, 32] --
│ └─Conv2d: 2-25 [12, 512, 32, 32] 1,180,160
│ └─BatchNorm2d: 2-26 [12, 512, 32, 32] 1,024
│ └─ReLU: 2-27 [12, 512, 32, 32] --
│ └─Conv2d: 2-28 [12, 512, 32, 32] 2,359,808
│ └─BatchNorm2d: 2-29 [12, 512, 32, 32] 1,024
│ └─ReLU: 2-30 [12, 512, 32, 32] --
│ └─Conv2d: 2-31 [12, 512, 32, 32] 2,359,808
│ └─BatchNorm2d: 2-32 [12, 512, 32, 32] 1,024
│ └─ReLU: 2-33 [12, 512, 32, 32] --
│ └─MaxPool2d: 2-34 [12, 512, 16, 16] --
│ └─Conv2d: 2-35 [12, 512, 16, 16] 2,359,808
│ └─BatchNorm2d: 2-36 [12, 512, 16, 16] 1,024
│ └─ReLU: 2-37 [12, 512, 16, 16] --
│ └─Conv2d: 2-38 [12, 512, 16, 16] 2,359,808
│ └─BatchNorm2d: 2-39 [12, 512, 16, 16] 1,024
│ └─ReLU: 2-40 [12, 512, 16, 16] --
│ └─Conv2d: 2-41 [12, 512, 16, 16] 2,359,808
│ └─BatchNorm2d: 2-42 [12, 512, 16, 16] 1,024
│ └─ReLU: 2-43 [12, 512, 16, 16] --
│ └─MaxPool2d: 2-44 [12, 512, 8, 8] --
├─AdaptiveAvgPool2d: 1-2 [12, 512, 7, 7] --
├─Sequential: 1-3 [12, 1000] --
│ └─Linear: 2-45 [12, 4096] 102,764,544
│ └─ReLU: 2-46 [12, 4096] --
│ └─Dropout: 2-47 [12, 4096] --
│ └─Linear: 2-48 [12, 4096] 16,781,312
│ └─ReLU: 2-49 [12, 4096] --
│ └─Dropout: 2-50 [12, 4096] --
│ └─Linear: 2-51 [12, 1000] 4,097,000
==========================================================================================
Total params: 138,365,992
Trainable params: 138,365,992
Non-trainable params: 0
Total mult-adds (Units.GIGABYTES): 242.23
==========================================================================================
Input size (MB): 9.44
Forward/backward pass size (MB): 3398.27
Params size (MB): 553.46
Estimated Total Size (MB): 3961.17
==========================================================================================It is possible to obtain a list of all of the architecture’s components using the features property of the instantiated VGGNet-16 architecture. Printing the result, you can see that the architecture has 43 components. I will need to divide the model’s components as follows to extract the results at the correct locations in the architecture.
- Layers 0-5 -> 1st skip connection (original array size)
- Layers 6-12 -> 2nd skip connection (original array size/2)
- Layers 13-22 -> 3rd skip connection (original array size/4)
- Layers 23-32 -> 4th skip connection (original array size/8)
- Layers 33-42 -> Bottleneck (original array size/16)
I will need to split the architecture before each max pooling operation so that I can match each encoder step with the associated decoder step that has the same array sizes in the spatial dimensions. The last max pooling layer is not used since I do not want do decrease the array size further. Instead, the data will enter the decoder component and the first 2D transpose convolution operation.
vgg16F = torchvision.models.vgg16_bn(weights=None).featuresvgg16FSequential(
(0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, bias=True, track_running_stats=True)
(2): ReLU(inplace=True)
(3): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(4): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, bias=True, track_running_stats=True)
(5): ReLU(inplace=True)
(6): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(7): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(8): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, bias=True, track_running_stats=True)
(9): ReLU(inplace=True)
(10): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(11): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, bias=True, track_running_stats=True)
(12): ReLU(inplace=True)
(13): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(14): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(15): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, bias=True, track_running_stats=True)
(16): ReLU(inplace=True)
(17): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(18): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, bias=True, track_running_stats=True)
(19): ReLU(inplace=True)
(20): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(21): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, bias=True, track_running_stats=True)
(22): ReLU(inplace=True)
(23): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(24): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(25): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, bias=True, track_running_stats=True)
(26): ReLU(inplace=True)
(27): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(28): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, bias=True, track_running_stats=True)
(29): ReLU(inplace=True)
(30): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(31): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, bias=True, track_running_stats=True)
(32): ReLU(inplace=True)
(33): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(34): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(35): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, bias=True, track_running_stats=True)
(36): ReLU(inplace=True)
(37): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(38): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, bias=True, track_running_stats=True)
(39): ReLU(inplace=True)
(40): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(41): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, bias=True, track_running_stats=True)
(42): ReLU(inplace=True)
(43): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
)With what we learned about the VGGNet-16 architecture from our explorations above, let’s now incorporate it into a UNet architecture. As I did in the first chapter focused on UNet, I define double_conv() and up_conv() functions for use in the UNet architecture. I next define the UNet architecture by subclassing nn.Module. Here are the key components.
- The
__init__()constructor method defines the following parameters: the number of input channels (nCls), the number of classes being differentiated (nCls), and whether or not to use pre-trained weights. - The features from the VGGNet-16 architecture are extracted using the
vgg16_bn()function from torchvision and the features property. - A list of output sizes for the decoder are defined as a list.
- The components of the encoder are defined as the appropriate subset of layers from VGGNet-16. The components are combined using
nn.Sequential(), and * is used to unpack the list of features. Again, I am breaking the model apart before each max pooling operation so that the array sizes in the spatial dimensions match between each encoder block and associated decoder block. - I define the bottleneck as the last set of components of the VGGNet-16 model. The last feature is not used since we do not want to apply the last max pooling operation.
- I define the decoder blocks, which each consist of upsampling using 2D transpose convolution and a series of two 2D convolution layers to learn filters.
- The last layer uses 2D convolution with a kernel size of 1x1 and a stride of 1. The number of outputs is equal to the number of classes being differentiated.
- The
forward()method defines how data will pass through the architecture. Data first pass through each of the encoder blocks followed by the bottleneck. In each decoder block, upsampling is performed using 2D transpose convolution, the feature maps from the associated encoder block are concatenated, and the features pass through two 2D convolution layers. Lastly, the data pass through the final 2D convolution layer, which has a kernel size of 1x1 and a stride of 1. The output will be logits for each class. In the case of a binary classification, it could return only a logit for the positive case. Probabilities are not returned since a sigmoid or softmax activation is not being applied.
def double_conv(inChannels, outChannels):
return nn.Sequential(
nn.Conv2d(inChannels, outChannels, kernel_size=(3,3), stride=1, padding=1),
nn.BatchNorm2d(outChannels),
nn.ReLU(inplace=True),
nn.Conv2d(outChannels, outChannels, kernel_size=(3,3), stride=1, padding=1),
nn.BatchNorm2d(outChannels),
nn.ReLU(inplace=True)
)def up_conv(inChannels, outChannels):
return nn.Sequential(
nn.ConvTranspose2d(inChannels, outChannels, kernel_size=(2,2), stride=2),
nn.BatchNorm2d(outChannels),
nn.ReLU(inplace=True)
)class myUNetVGG16(nn.Module):
def __init__(self, inChn, nCls, useWghts=True):
super().__init__()
self.inChn = inChn
self.nCls = nCls
self.useWghts = useWghts
self.base_model = torchvision.models.vgg16_bn(weights="DEFAULT" if useWghts else None).features
self.outSizes = [64, 128, 256, 512, 512]
self.encoder1 = nn.Sequential(*self.base_model[:6])
self.encoder2 = nn.Sequential(*self.base_model[6:13])
self.encoder3 = nn.Sequential(*self.base_model[13:23])
self.encoder4 = nn.Sequential(*self.base_model[23:33])
self.bottleneck = nn.Sequential(*self.base_model[33:43])
self.decoder1up = up_conv(self.outSizes[4], 512)
self.decoder1 = double_conv(self.outSizes[3] + 512, 256)
self.decoder2up = up_conv(256, 256)
self.decoder2 = double_conv(self.outSizes[2] + 256, 128)
self.decoder3up = up_conv(128, 128)
self.decoder3 = double_conv(self.outSizes[1] + 128, 64)
self.decoder4up = up_conv(64, 64)
self.decoder4 = double_conv(self.outSizes[0] + 64, 32)
self.classifier = nn.Conv2d(32, nCls, kernel_size=(1,1))
def forward(self, x):
#Encoder
encoder1 = self.encoder1(x)
encoder2 = self.encoder2(encoder1)
encoder3 = self.encoder3(encoder2)
encoder4 = self.encoder4(encoder3)
#Bottleneck
x = self.bottleneck(encoder4)
#Decoder
x = self.decoder1up(x)
x = torch.concat([x, encoder4], dim=1)
x = self.decoder1(x)
x = self.decoder2up(x)
x = torch.concat([x, encoder3], dim=1)
x = self.decoder2(x)
x = self.decoder3up(x)
x = torch.concat([x, encoder2], dim=1)
x = self.decoder3(x)
x = self.decoder4up(x)
x = torch.concat([x, encoder1], dim=1)
x = self.decoder4(x)
#Classifier head
x = self.classifier(x)
return xI instantiate an instance of the myUNetVGG16() architecture that accepts 3 input channels and outputs 10 class logits. I also initialize the model using the VGGNet-16 pre-trained weights available from torchvision.
Using the summary() function from torchinfo, you can see that the model has over 34 million trainable parameters. There are currently no non-trainable parameters. This is because, even though I downloaded the pre-trained weights, all parameters can still be updated during the learning process. In other words, the model will be initialized using these weights as opposed to random weights, but these layers and associated parameters are still trainable.
model = myUNetVGG16(inChn=3, nCls=10, useWghts=True).to(device)summary(model, (12,3,256,256))==========================================================================================
Layer (type:depth-idx) Output Shape Param #
==========================================================================================
myUNetVGG16 [12, 10, 256, 256] --
├─Sequential: 1-1 [12, 64, 256, 256] --
│ └─Conv2d: 2-1 [12, 64, 256, 256] 1,792
│ └─BatchNorm2d: 2-2 [12, 64, 256, 256] 128
│ └─ReLU: 2-3 [12, 64, 256, 256] --
│ └─Conv2d: 2-4 [12, 64, 256, 256] 36,928
│ └─BatchNorm2d: 2-5 [12, 64, 256, 256] 128
│ └─ReLU: 2-6 [12, 64, 256, 256] --
├─Sequential: 1-2 [12, 128, 128, 128] --
│ └─MaxPool2d: 2-7 [12, 64, 128, 128] --
│ └─Conv2d: 2-8 [12, 128, 128, 128] 73,856
│ └─BatchNorm2d: 2-9 [12, 128, 128, 128] 256
│ └─ReLU: 2-10 [12, 128, 128, 128] --
│ └─Conv2d: 2-11 [12, 128, 128, 128] 147,584
│ └─BatchNorm2d: 2-12 [12, 128, 128, 128] 256
│ └─ReLU: 2-13 [12, 128, 128, 128] --
├─Sequential: 1-3 [12, 256, 64, 64] --
│ └─MaxPool2d: 2-14 [12, 128, 64, 64] --
│ └─Conv2d: 2-15 [12, 256, 64, 64] 295,168
│ └─BatchNorm2d: 2-16 [12, 256, 64, 64] 512
│ └─ReLU: 2-17 [12, 256, 64, 64] --
│ └─Conv2d: 2-18 [12, 256, 64, 64] 590,080
│ └─BatchNorm2d: 2-19 [12, 256, 64, 64] 512
│ └─ReLU: 2-20 [12, 256, 64, 64] --
│ └─Conv2d: 2-21 [12, 256, 64, 64] 590,080
│ └─BatchNorm2d: 2-22 [12, 256, 64, 64] 512
│ └─ReLU: 2-23 [12, 256, 64, 64] --
├─Sequential: 1-4 [12, 512, 32, 32] --
│ └─MaxPool2d: 2-24 [12, 256, 32, 32] --
│ └─Conv2d: 2-25 [12, 512, 32, 32] 1,180,160
│ └─BatchNorm2d: 2-26 [12, 512, 32, 32] 1,024
│ └─ReLU: 2-27 [12, 512, 32, 32] --
│ └─Conv2d: 2-28 [12, 512, 32, 32] 2,359,808
│ └─BatchNorm2d: 2-29 [12, 512, 32, 32] 1,024
│ └─ReLU: 2-30 [12, 512, 32, 32] --
│ └─Conv2d: 2-31 [12, 512, 32, 32] 2,359,808
│ └─BatchNorm2d: 2-32 [12, 512, 32, 32] 1,024
│ └─ReLU: 2-33 [12, 512, 32, 32] --
├─Sequential: 1-5 [12, 512, 16, 16] --
│ └─MaxPool2d: 2-34 [12, 512, 16, 16] --
│ └─Conv2d: 2-35 [12, 512, 16, 16] 2,359,808
│ └─BatchNorm2d: 2-36 [12, 512, 16, 16] 1,024
│ └─ReLU: 2-37 [12, 512, 16, 16] --
│ └─Conv2d: 2-38 [12, 512, 16, 16] 2,359,808
│ └─BatchNorm2d: 2-39 [12, 512, 16, 16] 1,024
│ └─ReLU: 2-40 [12, 512, 16, 16] --
│ └─Conv2d: 2-41 [12, 512, 16, 16] 2,359,808
│ └─BatchNorm2d: 2-42 [12, 512, 16, 16] 1,024
│ └─ReLU: 2-43 [12, 512, 16, 16] --
├─Sequential: 1-6 [12, 512, 32, 32] --
│ └─ConvTranspose2d: 2-44 [12, 512, 32, 32] 1,049,088
│ └─BatchNorm2d: 2-45 [12, 512, 32, 32] 1,024
│ └─ReLU: 2-46 [12, 512, 32, 32] --
├─Sequential: 1-7 [12, 256, 32, 32] --
│ └─Conv2d: 2-47 [12, 256, 32, 32] 2,359,552
│ └─BatchNorm2d: 2-48 [12, 256, 32, 32] 512
│ └─ReLU: 2-49 [12, 256, 32, 32] --
│ └─Conv2d: 2-50 [12, 256, 32, 32] 590,080
│ └─BatchNorm2d: 2-51 [12, 256, 32, 32] 512
│ └─ReLU: 2-52 [12, 256, 32, 32] --
├─Sequential: 1-8 [12, 256, 64, 64] --
│ └─ConvTranspose2d: 2-53 [12, 256, 64, 64] 262,400
│ └─BatchNorm2d: 2-54 [12, 256, 64, 64] 512
│ └─ReLU: 2-55 [12, 256, 64, 64] --
├─Sequential: 1-9 [12, 128, 64, 64] --
│ └─Conv2d: 2-56 [12, 128, 64, 64] 589,952
│ └─BatchNorm2d: 2-57 [12, 128, 64, 64] 256
│ └─ReLU: 2-58 [12, 128, 64, 64] --
│ └─Conv2d: 2-59 [12, 128, 64, 64] 147,584
│ └─BatchNorm2d: 2-60 [12, 128, 64, 64] 256
│ └─ReLU: 2-61 [12, 128, 64, 64] --
├─Sequential: 1-10 [12, 128, 128, 128] --
│ └─ConvTranspose2d: 2-62 [12, 128, 128, 128] 65,664
│ └─BatchNorm2d: 2-63 [12, 128, 128, 128] 256
│ └─ReLU: 2-64 [12, 128, 128, 128] --
├─Sequential: 1-11 [12, 64, 128, 128] --
│ └─Conv2d: 2-65 [12, 64, 128, 128] 147,520
│ └─BatchNorm2d: 2-66 [12, 64, 128, 128] 128
│ └─ReLU: 2-67 [12, 64, 128, 128] --
│ └─Conv2d: 2-68 [12, 64, 128, 128] 36,928
│ └─BatchNorm2d: 2-69 [12, 64, 128, 128] 128
│ └─ReLU: 2-70 [12, 64, 128, 128] --
├─Sequential: 1-12 [12, 64, 256, 256] --
│ └─ConvTranspose2d: 2-71 [12, 64, 256, 256] 16,448
│ └─BatchNorm2d: 2-72 [12, 64, 256, 256] 128
│ └─ReLU: 2-73 [12, 64, 256, 256] --
├─Sequential: 1-13 [12, 32, 256, 256] --
│ └─Conv2d: 2-74 [12, 32, 256, 256] 36,896
│ └─BatchNorm2d: 2-75 [12, 32, 256, 256] 64
│ └─ReLU: 2-76 [12, 32, 256, 256] --
│ └─Conv2d: 2-77 [12, 32, 256, 256] 9,248
│ └─BatchNorm2d: 2-78 [12, 32, 256, 256] 64
│ └─ReLU: 2-79 [12, 32, 256, 256] --
├─Conv2d: 1-14 [12, 10, 256, 256] 330
==========================================================================================
Total params: 20,038,666
Trainable params: 20,038,666
Non-trainable params: 0
Total mult-adds (Units.GIGABYTES): 437.69
==========================================================================================
Input size (MB): 9.44
Forward/backward pass size (MB): 6480.20
Params size (MB): 80.15
Estimated Total Size (MB): 6569.79
==========================================================================================I next demonstrate using for loops to freeze the trainable parameters in the backbone or encoder layers. This is accomplished by iterating over all layers in the list of layers in the VGGNet-16 model followed by iterating over all of the parameters in each of these layers to set the requires_grad property to False.
If I print the summary again, you can now see that only a subset of the parameters is trainable. In other words, the parameters in the encoder component of the model that were defined using the VGGNet-16 model can no longer be updated. Only parameters in the decoder component will be trainable.
for l in model.base_model:
for param in l.parameters():
param.requires_grad = Falsesummary(model, (12,3,256,256))==========================================================================================
Layer (type:depth-idx) Output Shape Param #
==========================================================================================
myUNetVGG16 [12, 10, 256, 256] --
├─Sequential: 1-1 [12, 64, 256, 256] --
│ └─Conv2d: 2-1 [12, 64, 256, 256] (1,792)
│ └─BatchNorm2d: 2-2 [12, 64, 256, 256] (128)
│ └─ReLU: 2-3 [12, 64, 256, 256] --
│ └─Conv2d: 2-4 [12, 64, 256, 256] (36,928)
│ └─BatchNorm2d: 2-5 [12, 64, 256, 256] (128)
│ └─ReLU: 2-6 [12, 64, 256, 256] --
├─Sequential: 1-2 [12, 128, 128, 128] --
│ └─MaxPool2d: 2-7 [12, 64, 128, 128] --
│ └─Conv2d: 2-8 [12, 128, 128, 128] (73,856)
│ └─BatchNorm2d: 2-9 [12, 128, 128, 128] (256)
│ └─ReLU: 2-10 [12, 128, 128, 128] --
│ └─Conv2d: 2-11 [12, 128, 128, 128] (147,584)
│ └─BatchNorm2d: 2-12 [12, 128, 128, 128] (256)
│ └─ReLU: 2-13 [12, 128, 128, 128] --
├─Sequential: 1-3 [12, 256, 64, 64] --
│ └─MaxPool2d: 2-14 [12, 128, 64, 64] --
│ └─Conv2d: 2-15 [12, 256, 64, 64] (295,168)
│ └─BatchNorm2d: 2-16 [12, 256, 64, 64] (512)
│ └─ReLU: 2-17 [12, 256, 64, 64] --
│ └─Conv2d: 2-18 [12, 256, 64, 64] (590,080)
│ └─BatchNorm2d: 2-19 [12, 256, 64, 64] (512)
│ └─ReLU: 2-20 [12, 256, 64, 64] --
│ └─Conv2d: 2-21 [12, 256, 64, 64] (590,080)
│ └─BatchNorm2d: 2-22 [12, 256, 64, 64] (512)
│ └─ReLU: 2-23 [12, 256, 64, 64] --
├─Sequential: 1-4 [12, 512, 32, 32] --
│ └─MaxPool2d: 2-24 [12, 256, 32, 32] --
│ └─Conv2d: 2-25 [12, 512, 32, 32] (1,180,160)
│ └─BatchNorm2d: 2-26 [12, 512, 32, 32] (1,024)
│ └─ReLU: 2-27 [12, 512, 32, 32] --
│ └─Conv2d: 2-28 [12, 512, 32, 32] (2,359,808)
│ └─BatchNorm2d: 2-29 [12, 512, 32, 32] (1,024)
│ └─ReLU: 2-30 [12, 512, 32, 32] --
│ └─Conv2d: 2-31 [12, 512, 32, 32] (2,359,808)
│ └─BatchNorm2d: 2-32 [12, 512, 32, 32] (1,024)
│ └─ReLU: 2-33 [12, 512, 32, 32] --
├─Sequential: 1-5 [12, 512, 16, 16] --
│ └─MaxPool2d: 2-34 [12, 512, 16, 16] --
│ └─Conv2d: 2-35 [12, 512, 16, 16] (2,359,808)
│ └─BatchNorm2d: 2-36 [12, 512, 16, 16] (1,024)
│ └─ReLU: 2-37 [12, 512, 16, 16] --
│ └─Conv2d: 2-38 [12, 512, 16, 16] (2,359,808)
│ └─BatchNorm2d: 2-39 [12, 512, 16, 16] (1,024)
│ └─ReLU: 2-40 [12, 512, 16, 16] --
│ └─Conv2d: 2-41 [12, 512, 16, 16] (2,359,808)
│ └─BatchNorm2d: 2-42 [12, 512, 16, 16] (1,024)
│ └─ReLU: 2-43 [12, 512, 16, 16] --
├─Sequential: 1-6 [12, 512, 32, 32] --
│ └─ConvTranspose2d: 2-44 [12, 512, 32, 32] 1,049,088
│ └─BatchNorm2d: 2-45 [12, 512, 32, 32] 1,024
│ └─ReLU: 2-46 [12, 512, 32, 32] --
├─Sequential: 1-7 [12, 256, 32, 32] --
│ └─Conv2d: 2-47 [12, 256, 32, 32] 2,359,552
│ └─BatchNorm2d: 2-48 [12, 256, 32, 32] 512
│ └─ReLU: 2-49 [12, 256, 32, 32] --
│ └─Conv2d: 2-50 [12, 256, 32, 32] 590,080
│ └─BatchNorm2d: 2-51 [12, 256, 32, 32] 512
│ └─ReLU: 2-52 [12, 256, 32, 32] --
├─Sequential: 1-8 [12, 256, 64, 64] --
│ └─ConvTranspose2d: 2-53 [12, 256, 64, 64] 262,400
│ └─BatchNorm2d: 2-54 [12, 256, 64, 64] 512
│ └─ReLU: 2-55 [12, 256, 64, 64] --
├─Sequential: 1-9 [12, 128, 64, 64] --
│ └─Conv2d: 2-56 [12, 128, 64, 64] 589,952
│ └─BatchNorm2d: 2-57 [12, 128, 64, 64] 256
│ └─ReLU: 2-58 [12, 128, 64, 64] --
│ └─Conv2d: 2-59 [12, 128, 64, 64] 147,584
│ └─BatchNorm2d: 2-60 [12, 128, 64, 64] 256
│ └─ReLU: 2-61 [12, 128, 64, 64] --
├─Sequential: 1-10 [12, 128, 128, 128] --
│ └─ConvTranspose2d: 2-62 [12, 128, 128, 128] 65,664
│ └─BatchNorm2d: 2-63 [12, 128, 128, 128] 256
│ └─ReLU: 2-64 [12, 128, 128, 128] --
├─Sequential: 1-11 [12, 64, 128, 128] --
│ └─Conv2d: 2-65 [12, 64, 128, 128] 147,520
│ └─BatchNorm2d: 2-66 [12, 64, 128, 128] 128
│ └─ReLU: 2-67 [12, 64, 128, 128] --
│ └─Conv2d: 2-68 [12, 64, 128, 128] 36,928
│ └─BatchNorm2d: 2-69 [12, 64, 128, 128] 128
│ └─ReLU: 2-70 [12, 64, 128, 128] --
├─Sequential: 1-12 [12, 64, 256, 256] --
│ └─ConvTranspose2d: 2-71 [12, 64, 256, 256] 16,448
│ └─BatchNorm2d: 2-72 [12, 64, 256, 256] 128
│ └─ReLU: 2-73 [12, 64, 256, 256] --
├─Sequential: 1-13 [12, 32, 256, 256] --
│ └─Conv2d: 2-74 [12, 32, 256, 256] 36,896
│ └─BatchNorm2d: 2-75 [12, 32, 256, 256] 64
│ └─ReLU: 2-76 [12, 32, 256, 256] --
│ └─Conv2d: 2-77 [12, 32, 256, 256] 9,248
│ └─BatchNorm2d: 2-78 [12, 32, 256, 256] 64
│ └─ReLU: 2-79 [12, 32, 256, 256] --
├─Conv2d: 1-14 [12, 10, 256, 256] 330
==========================================================================================
Total params: 20,038,666
Trainable params: 5,315,530
Non-trainable params: 14,723,136
Total mult-adds (Units.GIGABYTES): 437.69
==========================================================================================
Input size (MB): 9.44
Forward/backward pass size (MB): 6480.20
Params size (MB): 80.15
Estimated Total Size (MB): 6569.79
==========================================================================================17.3 ResNet Encoder
A UNet with a ResNet-18 encoder is conceptualized in Figure 17.2. Since the first ResNet operation, the 7x7 convolution, has a stride of 2, an additional encoder block consisting of two 3x3 convolutional layers is applied to the input so that feature maps at the original spatial resolution are available to the decoder. The output of the 7x7 convolution and of each of the first three stages of residual blocks are then passed to the decoder through skip connections, and the output of the last stage serves as the input to the decoder.
As a second example, I will now define a UNet architecture that can accept different versions of a ResNet architecture as the backbone. I begin by redefining the double_conv() and up_conv() functions. In the myUNetResNet() subclass definition, I define the following parameters
- inChn = number of input channels
- inCls = number of output classes
- resNet = which ResNet architecture to use (“18”, “34”, or “50”); the default is “18”
- useWhgts = whether or not to initialize using pre-trained weights.
The backbone used will depend on the argument provided to the resNet parameter. Note that this will not change how the encoder blocks are defined since the list of layers does not change. The type or ReNet varies based on the number of operations in each layer, but the total number of layers remains the same. Similar to the UNet using the VGGNet-16 backbone, the layers from the ResNet must be partitioned into the appropriate encoder blocks such that the correct array size is extracted to concatenate with the associated decoder block. The first operation in the ResNet architecture actually reduces the array size, so the first encoder block uses the defined double_conv() operation as opposed to the first set of layers in the ResNet. Portions of the ResNet are used to define the operations in the next 4 encoder blocks and the bottleneck layer. Once these components are defined, the decoder components are defined, which consist of upsampling with 2D transpose convolution and learning new filters with a sequence of two 2D convolution layers.
In the forward method, as normal, how the data pass through the architecture is defined. The input data will first pass through the first encoder block. Since this is not part of the ResNet architecture, the original data are also passed through the second encoder block as opposed to the output from the first encoder block. The 2nd through 5th encoder blocks and the bottleneck block are defined based on components of the ResNet architecture.
Next, the data pass through the decoder component of the architecture. Each block consists of upsampling with 2D transpose convolution, concatenation of the feature maps from the associated encoder block, and using two 2D convolution layers to learn additional filters. Lastly, the data are then passed trough a 2D convolutional layer with a kernel size of 1x1 and a stride of 1 to obtain the class logits.
Again, it is important here that the correct stage of the ResNet architecture be assigned to the correct encoder block so that the skip connections will deliver arrays with the correct sizes in the spatial dimensions, which will then be concatenated with the layers from the bottleneck or prior decoder block. Note that there are other ways to integrate the ResNet backbone into a UNet architecture; this is just one example.
def double_conv(inChannels, outChannels):
return nn.Sequential(
nn.Conv2d(inChannels, outChannels, kernel_size=(3,3), stride=1, padding=1),
nn.BatchNorm2d(outChannels),
nn.ReLU(inplace=True),
nn.Conv2d(outChannels, outChannels, kernel_size=(3,3), stride=1, padding=1),
nn.BatchNorm2d(outChannels),
nn.ReLU(inplace=True)
)def up_conv(inChannels, outChannels):
return nn.Sequential(
nn.ConvTranspose2d(inChannels, outChannels, kernel_size=(2,2), stride=2),
nn.BatchNorm2d(outChannels),
nn.ReLU(inplace=True)
)class myUNetResNet(nn.Module):
def __init__(self, inChn, nCls, resNet = "18", useWghts=True):
super().__init__()
self.inChn = inChn
self.nCls = nCls
self.resNet = resNet
self.useWghts = useWghts
if(resNet == "34"):
self.base_model = torchvision.models.resnet34(weights="DEFAULT" if useWghts else None)
self.base_layers = list(self.base_model.children())
self.outSizes = [64, 64, 128, 256, 512]
elif(resNet == "50"):
self.base_model = torchvision.models.resnet50(weights="DEFAULT" if useWghts else None)
self.base_layers = list(self.base_model.children())
self.outSizes = [64, 256, 512, 1024, 2048]
else:
self.base_model = torchvision.models.resnet18(weights="DEFAULT" if useWghts else None)
self.base_layers = list(self.base_model.children())
self.outSizes = [64, 64, 128, 256, 512]
self.encoder1 = double_conv(inChn, 16)
self.encoder2 = nn.Sequential(*self.base_layers[:3])
self.encoder3 = nn.Sequential(*self.base_layers[3:5])
self.encoder4 = self.base_layers[5]
self.encoder5 = self.base_layers[6]
self.bottleneck = self.base_layers[7]
self.decoder1up = up_conv(self.outSizes[4], 512)
self.decoder1 = double_conv(self.outSizes[3] + 512, 256)
self.decoder2up = up_conv(256, 256)
self.decoder2 = double_conv(self.outSizes[2] + 256, 128)
self.decoder3up = up_conv(128, 128)
self.decoder3 = double_conv(self.outSizes[1] + 128, 64)
self.decoder4up = up_conv(64, 64)
self.decoder4 = double_conv(self.outSizes[0] + 64, 32)
self.decoder5up = up_conv(32, 32)
self.decoder5 = double_conv(16 + 32, 16)
self.classifier = nn.Conv2d(16, nCls, kernel_size=(1,1))
def forward(self, x):
#Encoder
encoder1 = self.encoder1(x)
encoder2 = self.encoder2(x)
encoder3 = self.encoder3(encoder2)
encoder4 = self.encoder4(encoder3)
encoder5 = self.encoder5(encoder4)
#Bottleneck
x = self.bottleneck(encoder5)
#Decoder
x = self.decoder1up(x)
x = torch.concat([x, encoder5], dim=1)
x = self.decoder1(x)
x = self.decoder2up(x)
x = torch.concat([x, encoder4], dim=1)
x = self.decoder2(x)
x = self.decoder3up(x)
x = torch.concat([x, encoder3], dim=1)
x = self.decoder3(x)
x = self.decoder4up(x)
x = torch.concat([x, encoder2], dim=1)
x = self.decoder4(x)
x = self.decoder5up(x)
x = torch.concat([x, encoder1], dim=1)
x = self.decoder5(x)
#Classifier head
x = self.classifier(x)
return xI instantiate an instance of the myUnetResNet subclass that accepts 3 channels, differentiates 10 classes, uses a ResNet-18 architecture in the backbone or encoder, and is initialized using pre-trained weights in the encoder. I print a summary to explore the model architecture, which has over 26 million trainable parameter.
model = myUNetResNet(inChn=3, nCls=10, resNet = "18", useWghts=True).to(device)summary(model, (12,3,256,256))==========================================================================================
Layer (type:depth-idx) Output Shape Param #
==========================================================================================
myUNetResNet [12, 10, 256, 256] 513,000
├─Sequential: 1-1 [12, 16, 256, 256] --
│ └─Conv2d: 2-1 [12, 16, 256, 256] 448
│ └─BatchNorm2d: 2-2 [12, 16, 256, 256] 32
│ └─ReLU: 2-3 [12, 16, 256, 256] --
│ └─Conv2d: 2-4 [12, 16, 256, 256] 2,320
│ └─BatchNorm2d: 2-5 [12, 16, 256, 256] 32
│ └─ReLU: 2-6 [12, 16, 256, 256] --
├─Sequential: 1-2 [12, 64, 128, 128] --
│ └─Conv2d: 2-7 [12, 64, 128, 128] 9,408
│ └─BatchNorm2d: 2-8 [12, 64, 128, 128] 128
│ └─ReLU: 2-9 [12, 64, 128, 128] --
├─Sequential: 1-3 [12, 64, 64, 64] --
│ └─MaxPool2d: 2-10 [12, 64, 64, 64] --
│ └─Sequential: 2-11 [12, 64, 64, 64] --
│ │ └─BasicBlock: 3-1 [12, 64, 64, 64] 73,984
│ │ └─BasicBlock: 3-2 [12, 64, 64, 64] 73,984
├─Sequential: 1-4 [12, 128, 32, 32] --
│ └─BasicBlock: 2-12 [12, 128, 32, 32] --
│ │ └─Conv2d: 3-3 [12, 128, 32, 32] 73,728
│ │ └─BatchNorm2d: 3-4 [12, 128, 32, 32] 256
│ │ └─ReLU: 3-5 [12, 128, 32, 32] --
│ │ └─Conv2d: 3-6 [12, 128, 32, 32] 147,456
│ │ └─BatchNorm2d: 3-7 [12, 128, 32, 32] 256
│ │ └─Sequential: 3-8 [12, 128, 32, 32] 8,448
│ │ └─ReLU: 3-9 [12, 128, 32, 32] --
│ └─BasicBlock: 2-13 [12, 128, 32, 32] --
│ │ └─Conv2d: 3-10 [12, 128, 32, 32] 147,456
│ │ └─BatchNorm2d: 3-11 [12, 128, 32, 32] 256
│ │ └─ReLU: 3-12 [12, 128, 32, 32] --
│ │ └─Conv2d: 3-13 [12, 128, 32, 32] 147,456
│ │ └─BatchNorm2d: 3-14 [12, 128, 32, 32] 256
│ │ └─ReLU: 3-15 [12, 128, 32, 32] --
├─Sequential: 1-5 [12, 256, 16, 16] --
│ └─BasicBlock: 2-14 [12, 256, 16, 16] --
│ │ └─Conv2d: 3-16 [12, 256, 16, 16] 294,912
│ │ └─BatchNorm2d: 3-17 [12, 256, 16, 16] 512
│ │ └─ReLU: 3-18 [12, 256, 16, 16] --
│ │ └─Conv2d: 3-19 [12, 256, 16, 16] 589,824
│ │ └─BatchNorm2d: 3-20 [12, 256, 16, 16] 512
│ │ └─Sequential: 3-21 [12, 256, 16, 16] 33,280
│ │ └─ReLU: 3-22 [12, 256, 16, 16] --
│ └─BasicBlock: 2-15 [12, 256, 16, 16] --
│ │ └─Conv2d: 3-23 [12, 256, 16, 16] 589,824
│ │ └─BatchNorm2d: 3-24 [12, 256, 16, 16] 512
│ │ └─ReLU: 3-25 [12, 256, 16, 16] --
│ │ └─Conv2d: 3-26 [12, 256, 16, 16] 589,824
│ │ └─BatchNorm2d: 3-27 [12, 256, 16, 16] 512
│ │ └─ReLU: 3-28 [12, 256, 16, 16] --
├─Sequential: 1-6 [12, 512, 8, 8] --
│ └─BasicBlock: 2-16 [12, 512, 8, 8] --
│ │ └─Conv2d: 3-29 [12, 512, 8, 8] 1,179,648
│ │ └─BatchNorm2d: 3-30 [12, 512, 8, 8] 1,024
│ │ └─ReLU: 3-31 [12, 512, 8, 8] --
│ │ └─Conv2d: 3-32 [12, 512, 8, 8] 2,359,296
│ │ └─BatchNorm2d: 3-33 [12, 512, 8, 8] 1,024
│ │ └─Sequential: 3-34 [12, 512, 8, 8] 132,096
│ │ └─ReLU: 3-35 [12, 512, 8, 8] --
│ └─BasicBlock: 2-17 [12, 512, 8, 8] --
│ │ └─Conv2d: 3-36 [12, 512, 8, 8] 2,359,296
│ │ └─BatchNorm2d: 3-37 [12, 512, 8, 8] 1,024
│ │ └─ReLU: 3-38 [12, 512, 8, 8] --
│ │ └─Conv2d: 3-39 [12, 512, 8, 8] 2,359,296
│ │ └─BatchNorm2d: 3-40 [12, 512, 8, 8] 1,024
│ │ └─ReLU: 3-41 [12, 512, 8, 8] --
├─Sequential: 1-7 [12, 512, 16, 16] --
│ └─ConvTranspose2d: 2-18 [12, 512, 16, 16] 1,049,088
│ └─BatchNorm2d: 2-19 [12, 512, 16, 16] 1,024
│ └─ReLU: 2-20 [12, 512, 16, 16] --
├─Sequential: 1-8 [12, 256, 16, 16] --
│ └─Conv2d: 2-21 [12, 256, 16, 16] 1,769,728
│ └─BatchNorm2d: 2-22 [12, 256, 16, 16] 512
│ └─ReLU: 2-23 [12, 256, 16, 16] --
│ └─Conv2d: 2-24 [12, 256, 16, 16] 590,080
│ └─BatchNorm2d: 2-25 [12, 256, 16, 16] 512
│ └─ReLU: 2-26 [12, 256, 16, 16] --
├─Sequential: 1-9 [12, 256, 32, 32] --
│ └─ConvTranspose2d: 2-27 [12, 256, 32, 32] 262,400
│ └─BatchNorm2d: 2-28 [12, 256, 32, 32] 512
│ └─ReLU: 2-29 [12, 256, 32, 32] --
├─Sequential: 1-10 [12, 128, 32, 32] --
│ └─Conv2d: 2-30 [12, 128, 32, 32] 442,496
│ └─BatchNorm2d: 2-31 [12, 128, 32, 32] 256
│ └─ReLU: 2-32 [12, 128, 32, 32] --
│ └─Conv2d: 2-33 [12, 128, 32, 32] 147,584
│ └─BatchNorm2d: 2-34 [12, 128, 32, 32] 256
│ └─ReLU: 2-35 [12, 128, 32, 32] --
├─Sequential: 1-11 [12, 128, 64, 64] --
│ └─ConvTranspose2d: 2-36 [12, 128, 64, 64] 65,664
│ └─BatchNorm2d: 2-37 [12, 128, 64, 64] 256
│ └─ReLU: 2-38 [12, 128, 64, 64] --
├─Sequential: 1-12 [12, 64, 64, 64] --
│ └─Conv2d: 2-39 [12, 64, 64, 64] 110,656
│ └─BatchNorm2d: 2-40 [12, 64, 64, 64] 128
│ └─ReLU: 2-41 [12, 64, 64, 64] --
│ └─Conv2d: 2-42 [12, 64, 64, 64] 36,928
│ └─BatchNorm2d: 2-43 [12, 64, 64, 64] 128
│ └─ReLU: 2-44 [12, 64, 64, 64] --
├─Sequential: 1-13 [12, 64, 128, 128] --
│ └─ConvTranspose2d: 2-45 [12, 64, 128, 128] 16,448
│ └─BatchNorm2d: 2-46 [12, 64, 128, 128] 128
│ └─ReLU: 2-47 [12, 64, 128, 128] --
├─Sequential: 1-14 [12, 32, 128, 128] --
│ └─Conv2d: 2-48 [12, 32, 128, 128] 36,896
│ └─BatchNorm2d: 2-49 [12, 32, 128, 128] 64
│ └─ReLU: 2-50 [12, 32, 128, 128] --
│ └─Conv2d: 2-51 [12, 32, 128, 128] 9,248
│ └─BatchNorm2d: 2-52 [12, 32, 128, 128] 64
│ └─ReLU: 2-53 [12, 32, 128, 128] --
├─Sequential: 1-15 [12, 32, 256, 256] --
│ └─ConvTranspose2d: 2-54 [12, 32, 256, 256] 4,128
│ └─BatchNorm2d: 2-55 [12, 32, 256, 256] 64
│ └─ReLU: 2-56 [12, 32, 256, 256] --
├─Sequential: 1-16 [12, 16, 256, 256] --
│ └─Conv2d: 2-57 [12, 16, 256, 256] 6,928
│ └─BatchNorm2d: 2-58 [12, 16, 256, 256] 32
│ └─ReLU: 2-59 [12, 16, 256, 256] --
│ └─Conv2d: 2-60 [12, 16, 256, 256] 2,320
│ └─BatchNorm2d: 2-61 [12, 16, 256, 256] 32
│ └─ReLU: 2-62 [12, 16, 256, 256] --
├─Conv2d: 1-17 [12, 10, 256, 256] 170
==========================================================================================
Total params: 16,247,074
Trainable params: 16,247,074
Non-trainable params: 0
Total mult-adds (Units.GIGABYTES): 84.99
==========================================================================================
Input size (MB): 9.44
Forward/backward pass size (MB): 2648.70
Params size (MB): 62.94
Estimated Total Size (MB): 2721.08
==========================================================================================Similar to the VGGNet-16 example above, I can freeze the backbone parameters to reduce the number of trainable parameters in the model. This is accomplished by setting the requires_grad property for the backbone or encoder layers extracted from the ResNet architecture to False. Printing the summary, you can see that only a subset of the parameters is now trainable.
for l in model.base_layers:
for param in l.parameters():
param.requires_grad = Falsesummary(model, (12,3,256,256))==========================================================================================
Layer (type:depth-idx) Output Shape Param #
==========================================================================================
myUNetResNet [12, 10, 256, 256] 513,000
├─Sequential: 1-1 [12, 16, 256, 256] --
│ └─Conv2d: 2-1 [12, 16, 256, 256] 448
│ └─BatchNorm2d: 2-2 [12, 16, 256, 256] 32
│ └─ReLU: 2-3 [12, 16, 256, 256] --
│ └─Conv2d: 2-4 [12, 16, 256, 256] 2,320
│ └─BatchNorm2d: 2-5 [12, 16, 256, 256] 32
│ └─ReLU: 2-6 [12, 16, 256, 256] --
├─Sequential: 1-2 [12, 64, 128, 128] --
│ └─Conv2d: 2-7 [12, 64, 128, 128] (9,408)
│ └─BatchNorm2d: 2-8 [12, 64, 128, 128] (128)
│ └─ReLU: 2-9 [12, 64, 128, 128] --
├─Sequential: 1-3 [12, 64, 64, 64] --
│ └─MaxPool2d: 2-10 [12, 64, 64, 64] --
│ └─Sequential: 2-11 [12, 64, 64, 64] --
│ │ └─BasicBlock: 3-1 [12, 64, 64, 64] (73,984)
│ │ └─BasicBlock: 3-2 [12, 64, 64, 64] (73,984)
├─Sequential: 1-4 [12, 128, 32, 32] --
│ └─BasicBlock: 2-12 [12, 128, 32, 32] --
│ │ └─Conv2d: 3-3 [12, 128, 32, 32] (73,728)
│ │ └─BatchNorm2d: 3-4 [12, 128, 32, 32] (256)
│ │ └─ReLU: 3-5 [12, 128, 32, 32] --
│ │ └─Conv2d: 3-6 [12, 128, 32, 32] (147,456)
│ │ └─BatchNorm2d: 3-7 [12, 128, 32, 32] (256)
│ │ └─Sequential: 3-8 [12, 128, 32, 32] (8,448)
│ │ └─ReLU: 3-9 [12, 128, 32, 32] --
│ └─BasicBlock: 2-13 [12, 128, 32, 32] --
│ │ └─Conv2d: 3-10 [12, 128, 32, 32] (147,456)
│ │ └─BatchNorm2d: 3-11 [12, 128, 32, 32] (256)
│ │ └─ReLU: 3-12 [12, 128, 32, 32] --
│ │ └─Conv2d: 3-13 [12, 128, 32, 32] (147,456)
│ │ └─BatchNorm2d: 3-14 [12, 128, 32, 32] (256)
│ │ └─ReLU: 3-15 [12, 128, 32, 32] --
├─Sequential: 1-5 [12, 256, 16, 16] --
│ └─BasicBlock: 2-14 [12, 256, 16, 16] --
│ │ └─Conv2d: 3-16 [12, 256, 16, 16] (294,912)
│ │ └─BatchNorm2d: 3-17 [12, 256, 16, 16] (512)
│ │ └─ReLU: 3-18 [12, 256, 16, 16] --
│ │ └─Conv2d: 3-19 [12, 256, 16, 16] (589,824)
│ │ └─BatchNorm2d: 3-20 [12, 256, 16, 16] (512)
│ │ └─Sequential: 3-21 [12, 256, 16, 16] (33,280)
│ │ └─ReLU: 3-22 [12, 256, 16, 16] --
│ └─BasicBlock: 2-15 [12, 256, 16, 16] --
│ │ └─Conv2d: 3-23 [12, 256, 16, 16] (589,824)
│ │ └─BatchNorm2d: 3-24 [12, 256, 16, 16] (512)
│ │ └─ReLU: 3-25 [12, 256, 16, 16] --
│ │ └─Conv2d: 3-26 [12, 256, 16, 16] (589,824)
│ │ └─BatchNorm2d: 3-27 [12, 256, 16, 16] (512)
│ │ └─ReLU: 3-28 [12, 256, 16, 16] --
├─Sequential: 1-6 [12, 512, 8, 8] --
│ └─BasicBlock: 2-16 [12, 512, 8, 8] --
│ │ └─Conv2d: 3-29 [12, 512, 8, 8] (1,179,648)
│ │ └─BatchNorm2d: 3-30 [12, 512, 8, 8] (1,024)
│ │ └─ReLU: 3-31 [12, 512, 8, 8] --
│ │ └─Conv2d: 3-32 [12, 512, 8, 8] (2,359,296)
│ │ └─BatchNorm2d: 3-33 [12, 512, 8, 8] (1,024)
│ │ └─Sequential: 3-34 [12, 512, 8, 8] (132,096)
│ │ └─ReLU: 3-35 [12, 512, 8, 8] --
│ └─BasicBlock: 2-17 [12, 512, 8, 8] --
│ │ └─Conv2d: 3-36 [12, 512, 8, 8] (2,359,296)
│ │ └─BatchNorm2d: 3-37 [12, 512, 8, 8] (1,024)
│ │ └─ReLU: 3-38 [12, 512, 8, 8] --
│ │ └─Conv2d: 3-39 [12, 512, 8, 8] (2,359,296)
│ │ └─BatchNorm2d: 3-40 [12, 512, 8, 8] (1,024)
│ │ └─ReLU: 3-41 [12, 512, 8, 8] --
├─Sequential: 1-7 [12, 512, 16, 16] --
│ └─ConvTranspose2d: 2-18 [12, 512, 16, 16] 1,049,088
│ └─BatchNorm2d: 2-19 [12, 512, 16, 16] 1,024
│ └─ReLU: 2-20 [12, 512, 16, 16] --
├─Sequential: 1-8 [12, 256, 16, 16] --
│ └─Conv2d: 2-21 [12, 256, 16, 16] 1,769,728
│ └─BatchNorm2d: 2-22 [12, 256, 16, 16] 512
│ └─ReLU: 2-23 [12, 256, 16, 16] --
│ └─Conv2d: 2-24 [12, 256, 16, 16] 590,080
│ └─BatchNorm2d: 2-25 [12, 256, 16, 16] 512
│ └─ReLU: 2-26 [12, 256, 16, 16] --
├─Sequential: 1-9 [12, 256, 32, 32] --
│ └─ConvTranspose2d: 2-27 [12, 256, 32, 32] 262,400
│ └─BatchNorm2d: 2-28 [12, 256, 32, 32] 512
│ └─ReLU: 2-29 [12, 256, 32, 32] --
├─Sequential: 1-10 [12, 128, 32, 32] --
│ └─Conv2d: 2-30 [12, 128, 32, 32] 442,496
│ └─BatchNorm2d: 2-31 [12, 128, 32, 32] 256
│ └─ReLU: 2-32 [12, 128, 32, 32] --
│ └─Conv2d: 2-33 [12, 128, 32, 32] 147,584
│ └─BatchNorm2d: 2-34 [12, 128, 32, 32] 256
│ └─ReLU: 2-35 [12, 128, 32, 32] --
├─Sequential: 1-11 [12, 128, 64, 64] --
│ └─ConvTranspose2d: 2-36 [12, 128, 64, 64] 65,664
│ └─BatchNorm2d: 2-37 [12, 128, 64, 64] 256
│ └─ReLU: 2-38 [12, 128, 64, 64] --
├─Sequential: 1-12 [12, 64, 64, 64] --
│ └─Conv2d: 2-39 [12, 64, 64, 64] 110,656
│ └─BatchNorm2d: 2-40 [12, 64, 64, 64] 128
│ └─ReLU: 2-41 [12, 64, 64, 64] --
│ └─Conv2d: 2-42 [12, 64, 64, 64] 36,928
│ └─BatchNorm2d: 2-43 [12, 64, 64, 64] 128
│ └─ReLU: 2-44 [12, 64, 64, 64] --
├─Sequential: 1-13 [12, 64, 128, 128] --
│ └─ConvTranspose2d: 2-45 [12, 64, 128, 128] 16,448
│ └─BatchNorm2d: 2-46 [12, 64, 128, 128] 128
│ └─ReLU: 2-47 [12, 64, 128, 128] --
├─Sequential: 1-14 [12, 32, 128, 128] --
│ └─Conv2d: 2-48 [12, 32, 128, 128] 36,896
│ └─BatchNorm2d: 2-49 [12, 32, 128, 128] 64
│ └─ReLU: 2-50 [12, 32, 128, 128] --
│ └─Conv2d: 2-51 [12, 32, 128, 128] 9,248
│ └─BatchNorm2d: 2-52 [12, 32, 128, 128] 64
│ └─ReLU: 2-53 [12, 32, 128, 128] --
├─Sequential: 1-15 [12, 32, 256, 256] --
│ └─ConvTranspose2d: 2-54 [12, 32, 256, 256] 4,128
│ └─BatchNorm2d: 2-55 [12, 32, 256, 256] 64
│ └─ReLU: 2-56 [12, 32, 256, 256] --
├─Sequential: 1-16 [12, 16, 256, 256] --
│ └─Conv2d: 2-57 [12, 16, 256, 256] 6,928
│ └─BatchNorm2d: 2-58 [12, 16, 256, 256] 32
│ └─ReLU: 2-59 [12, 16, 256, 256] --
│ └─Conv2d: 2-60 [12, 16, 256, 256] 2,320
│ └─BatchNorm2d: 2-61 [12, 16, 256, 256] 32
│ └─ReLU: 2-62 [12, 16, 256, 256] --
├─Conv2d: 1-17 [12, 10, 256, 256] 170
==========================================================================================
Total params: 16,247,074
Trainable params: 4,557,562
Non-trainable params: 11,689,512
Total mult-adds (Units.GIGABYTES): 84.99
==========================================================================================
Input size (MB): 9.44
Forward/backward pass size (MB): 2648.70
Params size (MB): 62.94
Estimated Total Size (MB): 2721.08
==========================================================================================17.4 MobileNetV2 Encoder
As a third example, I define a UNet that uses the MobileNetV2 architecture, which was introduced in Chapter 12, as the encoder. Since MobileNetV2 has far fewer parameters than VGGNet-16 or ResNet, the resulting UNet is much lighter, which reduces training time and memory use and speeds up inference over large geospatial extents. This makes MobileNetV2 a common choice when computational resources are limited or when a model must be applied to very large areas.
As with the other encoders, the key task is to identify where in the architecture to extract the feature maps that are passed to the decoder via the skip connections. In the torchvision implementation, the convolutional component of MobileNetV2 is stored in the features attribute, which contains 19 modules. The size of the spatial dimensions is halved by the modules at indices 0, 2, 4, 7, and 14, which use a stride of 2. So, I divide the modules into the following encoder blocks, each of which ends just before the next reduction in spatial dimensions:
- A custom double convolution block -> 1st skip connection (original size, 16 channels)
-
features[0:2]-> 2nd skip connection (original size/2, 16 channels) -
features[2:4]-> 3rd skip connection (original size/4, 24 channels) -
features[4:7]-> 4th skip connection (original size/8, 32 channels) -
features[7:14]-> 5th skip connection (original size/16, 96 channels) -
features[14:18]-> bottleneck (original size/32, 320 channels)
The final module, features[18], is a 1x1 convolution that expands the number of channels to 1,280 for the classification head of the original architecture. I do not include it since it greatly increases the number of parameters in the bottleneck without adding spatial context. As in the ResNet example, the first encoder block is a custom double convolution block applied to the input image, since the first MobileNetV2 module reduces the spatial dimensions and the decoder needs a skip connection at the original resolution. If the number of input channels is not 3, I replace the first convolution of MobileNetV2 so that it accepts the correct number of channels.
class myUNetMobileNet(nn.Module):
def __init__(self, inChn, nCls, useWghts=True):
super().__init__()
weights = "DEFAULT" if useWghts else None
base = torchvision.models.mobilenet_v2(weights=weights).features
if inChn != 3:
base[0][0] = nn.Conv2d(inChn, 32, kernel_size=3, stride=2, padding=1, bias=False)
self.encoder1 = double_conv(inChn, 16)
self.encoder2 = base[0:2] # H/2, 16 channels
self.encoder3 = base[2:4] # H/4, 24 channels
self.encoder4 = base[4:7] # H/8, 32 channels
self.encoder5 = base[7:14] # H/16, 96 channels
self.bottleneck = base[14:18] # H/32, 320 channels
self.encoderParts = nn.ModuleList([self.encoder2, self.encoder3, self.encoder4,
self.encoder5, self.bottleneck])
self.decoder1up = up_conv(320, 256)
self.decoder1 = double_conv(96 + 256, 128)
self.decoder2up = up_conv(128, 128)
self.decoder2 = double_conv(32 + 128, 64)
self.decoder3up = up_conv(64, 64)
self.decoder3 = double_conv(24 + 64, 32)
self.decoder4up = up_conv(32, 32)
self.decoder4 = double_conv(16 + 32, 16)
self.decoder5up = up_conv(16, 16)
self.decoder5 = double_conv(16 + 16, 16)
self.classifier = nn.Conv2d(16, nCls, kernel_size=(1,1))
def forward(self, x):
e1 = self.encoder1(x)
e2 = self.encoder2(x)
e3 = self.encoder3(e2)
e4 = self.encoder4(e3)
e5 = self.encoder5(e4)
x = self.bottleneck(e5)
x = self.decoder1(torch.cat([self.decoder1up(x), e5], dim=1))
x = self.decoder2(torch.cat([self.decoder2up(x), e4], dim=1))
x = self.decoder3(torch.cat([self.decoder3up(x), e3], dim=1))
x = self.decoder4(torch.cat([self.decoder4up(x), e2], dim=1))
x = self.decoder5(torch.cat([self.decoder5up(x), e1], dim=1))
return self.classifier(x)Since the bottleneck has a spatial resolution of 1/32 of the input, the height and width of the input image must be evenly divisible by 32. I instantiate the model with 3 input channels, 10 classes, and pre-trained ImageNet weights and summarize it. Compare the number of trainable parameters to the VGGNet-16 and ResNet-18 versions above.
model = myUNetMobileNet(inChn=3, nCls=10, useWghts=True).to(device)
summary(model, (12,3,256,256))As with the other encoders, the pre-trained encoder parameters can be frozen. Since I stored the encoder components in the encoderParts attribute, they can be frozen in a single loop. Note that if the number of input channels is not 3, you may want to leave the replaced first convolution trainable, since it is randomly initialized.
for param in model.encoderParts.parameters():
param.requires_grad = False
summary(model, (12,3,256,256))17.5 Model Summaries with PyTorch Lightning
Throughout this chapter, I have used torchinfo to summarize the architectures. As discussed in Chapter 10, PyTorch Lightning also provides model summaries, which are printed automatically at the start of training. Lightning’s summary is useful for confirming the number of trainable and non-trainable parameters, for example, to confirm that the encoder was frozen as intended. However, torchinfo provides a more detailed, layer-by-layer view, including the output shape of every layer, which is more useful when building an architecture and checking that the skip connections deliver tensors with the expected shapes. For this reason, I continue to use torchinfo when designing architectures. The example below summarizes the MobileNetV2 UNet using Lightning after wrapping it in the LitSegmenter LightningModule defined in Chapter 16. The class is redefined below so that this chapter’s code can be run on its own.
import lightning as L
import torchmetrics as tm
class LitSegmenter(L.LightningModule):
def __init__(self, model, nCls, criterion, lr=1e-4, maxLR=1e-2):
super().__init__()
self.save_hyperparameters(ignore=["model", "criterion"])
self.model = model
self.criterion = criterion
self.lr = lr
self.maxLR = maxLR
metrics = tm.MetricCollection({
"acc": tm.Accuracy(task="multiclass", num_classes=nCls, average="micro"),
"f1": tm.F1Score(task="multiclass", num_classes=nCls, average="macro"),
"kappa": tm.CohenKappa(task="multiclass", num_classes=nCls),
})
self.train_metrics = metrics.clone(prefix="train_")
self.val_metrics = metrics.clone(prefix="val_")
self.test_metrics = metrics.clone(prefix="test_")
def forward(self, x):
return self.model(x)
def _shared_step(self, batch, metrics, stage):
x, y = batch
logits = self(x)
loss = self.criterion(logits, y)
metrics.update(logits, y)
self.log(f"{stage}_loss", loss, on_step=False, on_epoch=True, prog_bar=True)
self.log_dict(metrics, on_step=False, on_epoch=True, prog_bar=(stage == "val"))
return loss
def training_step(self, batch, batch_idx):
return self._shared_step(batch, self.train_metrics, "train")
def validation_step(self, batch, batch_idx):
self._shared_step(batch, self.val_metrics, "val")
def test_step(self, batch, batch_idx):
self._shared_step(batch, self.test_metrics, "test")
def predict_step(self, batch, batch_idx):
x = batch[0] if isinstance(batch, (list, tuple)) else batch
return torch.softmax(self(x), dim=1)
def configure_optimizers(self):
optimizer = torch.optim.AdamW(self.parameters(), lr=self.lr)
scheduler = torch.optim.lr_scheduler.OneCycleLR(
optimizer, max_lr=self.maxLR, three_phase=True,
total_steps=self.trainer.estimated_stepping_batches)
return {"optimizer": optimizer,
"lr_scheduler": {"scheduler": scheduler, "interval": "step"}}from lightning.pytorch.utilities.model_summary import ModelSummary
from kornia import losses
litModel = LitSegmenter(model=model, nCls=10, criterion=losses.DiceLoss(average="macro")).to(device)
litModel.example_input_array = torch.zeros(12, 3, 256, 256)
print(ModelSummary(litModel, max_depth=2))17.6 Concluding Remarks
You can now define a basic UNet architecture, train a UNet model, and define UNet architectures that make use of common CNN architectures as the backbone or encoder and can accept pre-trained weights. In the next chapter we will explore modules that can be added to UNet to potentially improve its performance. We will then use the UNet trained in Chapter 16 to make predictions to an entire image extent (Chapter 19), and in Chapter 20 we will generate the image chips and masks used to train the models. There are other semantic segmentation architectures that are more complex and difficult to build from scratch. Also, you may want to be able to use a wide variety of backbones in a variety of different semantic segmentation architectures. In Chapter 21, we will explore the Segmentation Models package, which builds on PyTorch and allows for using many different semantic segmentation architectures, backbone encoders, and pre-trained weights without having to build them or define the model or components on your own.
17.7 Questions
- What is the role of the encoder or backbone in a semantic segmentation architecture? Why can CNN architectures designed for scene labeling, such as VGGNet-16 and ResNet, be used as encoders?
- Why must the VGGNet-16 architecture be split before each max pooling operation when it is used as the encoder for a UNet? How were the locations used to split MobileNetV2 into encoder blocks determined, and why must the input height and width be divisible by 32 for the MobileNetV2 UNet?
- Why is the final max pooling operation from VGGNet-16 not used in the UNet architecture?
- Why do the UNets with ResNet and MobileNetV2 encoders use a custom double convolution block, as opposed to the first layers of the pre-trained architecture, as their first encoder block?
- Why must the spatial dimensions of the feature maps passed through each skip connection match those of the associated decoder block?
- Explain how the encoder parameters were frozen in this chapter and how this impacted the number of trainable parameters.
- Explain the difference between initializing an encoder using pre-trained weights and freezing the encoder.
- What are the potential benefits of using a pre-trained encoder when only a small training dataset is available? What are the potential drawbacks?