DeepLabV3 architecture in medical image analysis

When traditional convolutional neural networks are used to extract features, spatial information is often lost during the sampling and pooling process. However, in the field of biomedical segmentation, this information is very important. Moreover, a very good resolution is necessary for the network to be able to discriminate the different anatomical parts (some of them can be of subcentimetric size). In order to address this problem, the DeepLab family of convolutional neural networks suggests extending the receptive field of convolutions in order to obtain more contextual information and reduce the loss of localization information. [1]

In this blog post, we review and detail the DeepLabV3 architecture in the context of biomedical image segmentation.

DeepLabV3 network architecture

The DeepLabV3 architecture is a powerful convolutional neural network for semantic image segmentation. Its main strength lies in its ability to capture information at different scales. [2]
The features extracted by the model represent either individual measurable properties or a phenomenon. They are discriminating and their choice is crucial for the efficiency of pattern recognition algorithms. We know that convolutional neural networks see the input image through a smaller window (called convolution) that moves over it to build its feature map. For example, a first feature map looks for curves, a second for a combination of curves that form circles, a third for a bicycle from lines and circles, and so on [3]. In the case of convolutional neural networks, these features are thumbnails of smaller resolution and depend on the convolutional filtering used. However, in semantic segmentation of medical images, the decrease in resolution can be critical for the segmentation of subcentimeter anatomical structures.

More specifically, the DeepLabV3 architecture consists of a backbone network which, from the input image, extracts the feature map on which the rest of the network is based. In biomedical image segmentation, the latter must have a high resolution in order to preserve the global information of the image.
In addition to the backbone network, the DeepLabV3 architecture has a main module for Atrous Spatial Pyramid Pooling (ASPP). This module captures information at various levels by relying on the principle of dilated (or atrous) convolution, which we will describe in detail below. This approach allows to keep features localized in order to obtain multi-scale context information and improve the segmentation process. [4]

The DeepLabV3 architecture proposes to use a pyramid of dilated convolutions to increase the receptive field of these convolutions without increasing the size of the convolution kernel or decreasing the resolution of the images from the backbone network. Indeed, this pyramid integrates the pointwise convolution mechanism to reduce computation time.

Presentation of the pointwise convolution

Pointwise convolution is a form of convolution using a kernel of size 1x1 iterating through the channels of the image for each pixel of the image. This kernel has a depth equal to the number of channels in the input image. It can be used:

• in conjunction with depthwise convolutions to create depthwise separable convolutions, useful for applying multiple filters with a reduced number of operations; [6]
• to change the dimensionality of the input (i.e., the number of channels); for example, the Inception architecture concatenates features computed from separate kernels resulting in an excessive number of channels, pointwise convolution is employed to control this.

Application of the pointwise convolution

To illustrate how pointwise convolution works, let's take an example of how it is applied on an image with the dimensions 12x12x3. First, a depth convolution is applied to the input image. This filters each channel of the image to preserve its depth. That is to say that by using three kernels of dimension 5x5x1, the image of dimension 12x12x3 is transformed into an image of dimension 8x8x3. Indeed, each core of dimension 5x5x1 flows through a channel of the image and gives an image of dimension 8x8x1. An image of dimension 8x8x3 is obtained by stacking these images. Thus, to generate an 8x8x1 image, we apply a kernel of dimension 1x1x3 to the previously obtained image of dimension 8x8x3. The pointwise convolution transforms a 3 channel image into a single channel image. Now, if the number of channels in the image is to be increased from a dimension 3 to a dimension 256, then 256 kernels of dimension 1x1x3 are constructed, each individually producing an image of dimension 8x8x1.

While a standard convolution would have led to a number of operations equal to 5x5x3x8x256, that is 1 228 800 operations, the pointwise convolution only reaches: 5x5x3x8x8 + 8x8x3x256, that is 4800 + 49152 = 53952 operations. [6]

Atrous convolution

The pointwise convolution allows to obtain a dilated convolution (or "atrous convolution") by adding a dilation rate r, which defines the spacing between the cells of the kernel. This allows to increase the size of the kernel, and thus the receiving field, without having to process too many pixels. This results in considerable speed gains without loss of accuracy.

Formally, the dilated convolution is written :

with r the expansion rate and K the size of the kernel applied to the input x [5]. We notice that if r is equal to 1, then we obtain a simple convolution.

If r is greater than 1, then spaces are introduced between the elements of the kernel. The parameter r defines the enlargement of the filter compared to the size of the simple filter (r = 1) that can progressively process a larger part of the image. Thus, the use of dilated convolutions in the architecture allows, without compressing the input too much, not to lose the localization information as it is the case with classical CNN architectures. [4]

Presentation of the Atrous Spatial Pyramid Pooling module

The Atrous Spatial Pyramid Pooling (ASPP) module is inspired by spatial pyramid pooling [3]. With different dilation rates, it efficiently captures multi-scale information. Its convolutions applied in parallel on the feature map form a spatial pyramid where the convolution layers expanded at different rates are then concatenated. Specifically, the ASPP module consists of:

• a 1×1 convolution to keep the feature map information useful for the network to discriminate some structures; 
• of three 3×3 convolutions in parallel with expansion rates respectively equal to 6, 12, 18, and an output step equal to 16 allowing to capture the context at different scales, these values are the best found empirically in the paper [2]; 
• subsampling of the map to smooth out contours and highlight broad features;
• concatenation of the resulting feature maps; then 
• a 1×1 convolution to reduce the number of feature cards to 256; 
• finally, an oversampling of the size of the feature map to that of the final image. [4]

Example of implementation

Implementation of the ASPP module [9]

The output of the backbone network that extracts the feature map imposes 512 channels for it. This number of channels is passed through dilated convolutions. Each of them includes an addition of pixels around the image in order not to decrease the size of the image. After each convolution a normalization is applied on a given number of input images so that the model adjusts its weights with respect to an error regrouping several input images. Then, the ReLU activation function will modify the data in a nonlinear way in order to have a new approach on these data. The information passes or not if the stimulation threshold of the activation function is reached. This one corresponds to a numerical limit, if the information transcribed by a number is lower than the threshold then the output value is 0 otherwise it is equal to 1.

class ASPP(nn.Module):
     def __init__(self, num_channel):
          super(ASPP, self).__init__()
         # 512 corresponds to the number of filters at the output of 
the Resnet and 256 at the output of the ASPP module
         self.conv_1x1_1 = nn.Conv3d(512, 256, kernel_size=1)
         self.bn_conv_1x1_1 = nn.BatchNorm3d(256)

         self.conv_3x3_1 = nn.Conv3d(512, 256, kernel_size=3, 
stride=1, padding=6, dilation=6)
         self.bn_conv_3x3_1 = nn.BatchNorm3d(256)

        self.conv_3x3_2 = nn.Conv3d(512, 256, kernel_size=3, 
stride=1, padding=12, dilation=12)
        self.bn_conv_3x3_2 = nn.BatchNorm3d(256)

        self.conv_3x3_3 = nn.Conv3d(512, 256, kernel_size=3, 
stride=1, padding=18, dilation=18)
        self.bn_conv_3x3_3 = nn.BatchNorm3d(256)

        self.avg_pool = nn.AdaptiveAvgPool3d(1)

        self.conv_1x1_2 = nn.Conv3d(512, 256, kernel_size=1)
        self.bn_conv_1x1_2 = nn.BatchNorm3d(256)

       # 1280 corresponds to the concatenation of the convolutions 
applied in parallel

       self.conv_1x1_3 = nn.Conv3d(1280, 256, kernel_size=1)
       self.bn_conv_1x1_3 = nn.BatchNorm3d(256)

       self.conv_1x1_4 = nn.Conv3d(256, num_channel, kernel_size=1)

The pyramid consists of 5 convolutions of 256 channels each. Thus, after concatenation, the feature map has 1280 channels. Applying an additional point convolution results in a feature map of 256 channels.

def forward(self, feature_map):
           feature_map_h = feature_map.size()[2]
           feature_map_w = feature_map.size()[3]
           feature_map_c = feature_map.size()[4]

           # Convolution
           out_1x1 =
F.relu(self.bn_conv_1x1_1(self.conv_1x1_1(feature_map)))
           # Convolution dilated
           out_3x3_1 =
F.relu(self.bn_conv_3x3_1(self.conv_3x3_1(feature_map)))
           out_3x3_2 =
F.relu(self.bn_conv_3x3_2(self.conv_3x3_2(feature_map)))
           out_3x3_3 =
F.relu(self.bn_conv_3x3_3(self.conv_3x3_3(feature_map)))
           # Sharing
           out_img = self.avg_pool(feature_map)
           out_img = 
F.relu(self.bn_conv_1x1_2(self.conv_1x1_2(out_img))
           # oversampling to find the size of the feature map
           out_img = F.interpolate(out_img, size=(feature_map_h,
           feature_map_w, feature_map_c), mode='trilinear', align_corners=True)

           # Concatenation of the 5 operations performed in parallel
           out = torch.cat([out_1x1, out_3x3_1, out_3x3_2, out_3x3_3, 
out_img], 1)
           out = F.relu(self.bn_conv_1x1_3(self.conv_1x1_3(out))

           # change from 256 filters to the number of channels in the 
input image
           out = self.conv_1x1_4(out)

           return out

Propagation of the input image in the network

First, the image feature map is extracted by a ResNet-34 backbone network. It has a resolution 16 times smaller than the input image. This architecture retains the location information through links between the previous layers. Then, this feature map is passed to the ASPP module to obtain a multi-scale information of it. Finally, to obtain an output image size similar to that of the input image an oversampling is used.

def forward(self, x):

          h = x.size()[2]
          w = x.size()[3]
          c = x.size()[4]

          feature_map = self.resnet(x)
          output = self.aspp(feature_map)
          output = F.interpolate(output, size=(h, w, c), mode='trilinear', align_corners=True)

         if self.last_activation.lower() == 'sigmoid':
              output = nn.Sigmoid()(output)
         elif self.last_activation.lower() == 'softmax':
              output = nn.Softmax()(output)

         return output

The last activation function of the DeepLabV3 network transforms the information contained in the feature maps into a probability of belonging to a class. The use of a softmax activation function is interesting in particular for a multi-class segmentation, since it distributes in probability the belonging to each class for a given pixel (the sum of these probabilities gives 1 for a given pixel). Thus, the highest probability value is the class  of the pixel. For a binary segmentation, the use of a sigmoid activation function is preferable to obtain a probability value between 0 and 1. Then, the membership of the pixel to class 1 (respectively 0) is assigned when the probability is higher (respectively lower) than 0.5.

Conclusion

The DeepLabV3 architecture jointly uses the dilated convolution and spatial pyramid pooling mechanisms to extract feature maps with a large context and encode multi-scale information. Indeed, the Atrous Spatial Pyramid Pooling (ASPP) module progressively doubles the dilation rates to extract features at multiple levels. The difference in resource consumption is noticeable with other pyramid modules using convolutions with kernels of different sizes. Here, the addition of dilated convolutions allows to enlarge the field of view without increasing the size of the filtering kernel, therefore the computation time. Thus, the DeepLabV3 architecture is suitable for the recognition of structures appearing at various levels, as it can be the case for biomedical images. However, the reproduction of shape contours remains difficult. This is why DeeplabV3 has been improved with a decoding module: DeeplabV3+. This version has added the missing information at a higher resolution, allowing to reconstruct the image with a better accuracy on the edges. We will discuss this architecture in more detail in a future article.