Deeplab v3+的结构代码简要分析

添加了解码模块来重构精确的图像物体边界。对比如图

deeplab v3+采用了与deeplab v3类似的多尺度带洞卷积结构ASPP，然后通过上采样，以及与不同卷积层相拼接，最终经过卷积以及上采样得到结果。

deeplab v3:

基于提出的编码-解码结构，可以任意通过控制 atrous convolution 来输出编码特征的分辨率，来平衡精度和运行时间(已有编码-解码结构不具有该能力.).

可以用来挖掘不同尺度的上下文信息

PSPNet 对不同尺度的网络进行池化处理，处理多尺度的上下文内容信息

deeplab v3+以resnet101为backbone

 import math
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 import torch.utils.model_zoo as model_zoo
 from modeling.sync_batchnorm.batchnorm import SynchronizedBatchNorm2d

 BatchNorm2d = SynchronizedBatchNorm2d

 class Bottleneck(nn.Module):
　　　　　　#'resnet网络的基本框架’
     expansion = 

     def __init__(self, inplanes, planes, stride=, dilation=, downsample=None):
         super(Bottleneck, self).__init__()
         self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=, bias=False)
         self.bn1 = BatchNorm2d(planes)
         self.conv2 = nn.Conv2d(planes, planes, kernel_size=, stride=stride,
                                dilation=dilation, padding=dilation, bias=False)
         self.bn2 = BatchNorm2d(planes)
         self.conv3 = nn.Conv2d(planes, planes * , kernel_size=, bias=False)
         self.bn3 = BatchNorm2d(planes * )
         self.relu = nn.ReLU(inplace=True)
         self.downsample = downsample
         self.stride = stride
         self.dilation = dilation

     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)
         out = self.relu(out)

         out = self.conv3(out)
         out = self.bn3(out)

         if self.downsample is not None:
             residual = self.downsample(x)

         out += residual
         out = self.relu(out)

         return out

 class ResNet(nn.Module):
 　　#renet网络的构成部分
     def __init__(self, nInputChannels, block, layers, os=, pretrained=False):
         self.inplanes =
         super(ResNet, self).__init__()
         if os == :
             strides = [, , , ]
             dilations = [, , , ]
             blocks = [, , ]
         elif os == :
             strides = [, , , ]
             dilations = [, , , ]
             blocks = [, , ]
         else:
             raise NotImplementedError

         # Modules
         self.conv1 = nn.Conv2d(nInputChannels, , kernel_size=, stride=, padding=,
                                 bias=False)
         self.bn1 = BatchNorm2d()
         self.relu = nn.ReLU(inplace=True)
         self.maxpool = nn.MaxPool2d(kernel_size=, stride=, padding=)

         self.layer1 = self._make_layer(block, , layers[], stride=strides[], dilation=dilations[])
         self.layer2 = self._make_layer(block, , layers[], stride=strides[], dilation=dilations[])
         self.layer3 = self._make_layer(block, , layers[], stride=strides[], dilation=dilations[])
         self.layer4 = self._make_MG_unit(block, , blocks=blocks, stride=strides[], dilation=dilations[])

         self._init_weight()

         if pretrained:
             self._load_pretrained_model()

     def _make_layer(self, block, planes, blocks, stride=, dilation=):
         downsample = None
         if stride !=  or self.inplanes != planes * block.expansion:
             downsample = nn.Sequential(
                 nn.Conv2d(self.inplanes, planes * block.expansion,
                           kernel_size=, stride=stride, bias=False),
                 BatchNorm2d(planes * block.expansion),
             )

         layers = []
         layers.append(block(self.inplanes, planes, stride, dilation, downsample))
         self.inplanes = planes * block.expansion
         for i in range(, blocks):
             layers.append(block(self.inplanes, planes))

         return nn.Sequential(*layers)

     def _make_MG_unit(self, block, planes, blocks=[, , ], stride=, dilation=):
         downsample = None
         if stride !=  or self.inplanes != planes * block.expansion:
             downsample = nn.Sequential(
                 nn.Conv2d(self.inplanes, planes * block.expansion,
                           kernel_size=, stride=stride, bias=False),
                 BatchNorm2d(planes * block.expansion),
             )

         layers = []
         layers.append(block(self.inplanes, planes, stride, dilation=blocks[]*dilation, downsample=downsample))
         self.inplanes = planes * block.expansion
         for i in range(, len(blocks)):
             layers.append(block(self.inplanes, planes, stride=, dilation=blocks[i]*dilation))

         return nn.Sequential(*layers)

     def forward(self, input):
         x = self.conv1(input)
         x = self.bn1(x)
         x = self.relu(x)
         x = self.maxpool(x)

         x = self.layer1(x)
         low_level_feat = x
         x = self.layer2(x)
         x = self.layer3(x)
         x = self.layer4(x)
         return x, low_level_feat

     def _init_weight(self):
         for m in self.modules():
             if isinstance(m, nn.Conv2d):
                 n = m.kernel_size[] * m.kernel_size[] * m.out_channels
                 m.weight.data.normal_(, math.sqrt(. / n))
             elif isinstance(m, BatchNorm2d):
                 m.weight.data.fill_()
                 m.bias.data.zero_()

     def _load_pretrained_model(self):
         pretrain_dict = model_zoo.load_url('https://download.pytorch.org/models/resnet101-5d3b4d8f.pth')
         model_dict = {}
         state_dict = self.state_dict()
         for k, v in pretrain_dict.items():
             if k in state_dict:
                 model_dict[k] = v
         state_dict.update(model_dict)
         self.load_state_dict(state_dict)

 def ResNet101(nInputChannels=, os=, pretrained=False):
     model = ResNet(nInputChannels, Bottleneck, [, , , ], os, pretrained=pretrained)
     return model

 class ASPP_module(nn.Module):
 　　#ASpp模块的组成
     def __init__(self, inplanes, planes, dilation):
         super(ASPP_module, self).__init__()
         if dilation == :
             kernel_size =
             padding =
         else:
             kernel_size =
             padding = dilation
         self.atrous_convolution = nn.Conv2d(inplanes, planes, kernel_size=kernel_size,
                                             stride=, padding=padding, dilation=dilation, bias=False)
         self.bn = BatchNorm2d(planes)
         self.relu = nn.ReLU()

         self._init_weight()

     def forward(self, x):
         x = self.atrous_convolution(x)
         x = self.bn(x)

         return self.relu(x)

     def _init_weight(self):
         for m in self.modules():
             if isinstance(m, nn.Conv2d):
                 n = m.kernel_size[] * m.kernel_size[] * m.out_channels
                 m.weight.data.normal_(, math.sqrt(. / n))
             elif isinstance(m, BatchNorm2d):
                 m.weight.data.fill_()
                 m.bias.data.zero_()

 class DeepLabv3_plus(nn.Module):
　　#正式开始deeplabv3+的结构组成
     def __init__(self, nInputChannels=, n_classes=, os=, pretrained=False, freeze_bn=False, _print=True):
         if _print:
             print("Constructing DeepLabv3+ model...")
             print("Backbone: Resnet-101")
             print("Number of classes: {}".format(n_classes))
             print("Output stride: {}".format(os))
             print("Number of Input Channels: {}".format(nInputChannels))
         super(DeepLabv3_plus, self).__init__()

         # Atrous Conv  首先获得从resnet101中提取的features map
         self.resnet_features = ResNet101(nInputChannels, os, pretrained=pretrained)

         # ASPP,挑选参数
         if os == :
             dilations = [, , , ]
         elif os == :
             dilations = [, , , ]
         else:
             raise NotImplementedError
 　　　　#四个不同带洞卷积的设置，获取不同感受野
         self.aspp1 = ASPP_module(, , dilation=dilations[])
         self.aspp2 = ASPP_module(, , dilation=dilations[])
         self.aspp3 = ASPP_module(, , dilation=dilations[])
         self.aspp4 = ASPP_module(, , dilation=dilations[])

         self.relu = nn.ReLU()
 　　　　#全局平均池化层的设置
         self.global_avg_pool = nn.Sequential(nn.AdaptiveAvgPool2d((, )),
                                              nn.Conv2d(, , , stride=, bias=False),
                                              BatchNorm2d(),
                                              nn.ReLU())

         self.conv1 = nn.Conv2d(, , , bias=False)
         self.bn1 = BatchNorm2d()

         # adopt [1x1, ] for channel reduction.
         self.conv2 = nn.Conv2d(, , , bias=False)
         self.bn2 = BatchNorm2d()
 　　　　#结构图中的解码部分的最后一个3*3的卷积块
         self.last_conv = nn.Sequential(nn.Conv2d(, , kernel_size=, stride=, padding=, bias=False),
                                        BatchNorm2d(),
                                        nn.ReLU(),
                                        nn.Conv2d(, , kernel_size=, stride=, padding=, bias=False),
                                        BatchNorm2d(),
                                        nn.ReLU(),
                                        nn.Conv2d(, n_classes, kernel_size=, stride=))
         if freeze_bn:
             self._freeze_bn()
 　　#前向传播
     def forward(self, input):
         x, low_level_features = self.resnet_features(input)
         x1 = self.aspp1(x)
         x2 = self.aspp2(x)
         x3 = self.aspp3(x)
         x4 = self.aspp4(x)
         x5 = self.global_avg_pool(x)
         x5 = F.upsample(x5, size=x4.size()[:], mode='bilinear', align_corners=True)
 　　　　#把四个ASPP模块以及全局池化层拼接起来
         x = torch.cat((x1, x2, x3, x4, x5), dim=)
 　　　　#上采样
         x = self.conv1(x)
         x = self.bn1(x)
         x = self.relu(x)
         x = F.upsample(x, size=(int(math.ceil(input.size()[-]/)),
                                 int(math.ceil(input.size()[-]/))), mode='bilinear', align_corners=True)

         low_level_features = self.conv2(low_level_features)
         low_level_features = self.bn2(low_level_features)
         low_level_features = self.relu(low_level_features)

 　　　　　#拼接低层次的特征，然后再通过插值获取原图大小的结果
         x = torch.cat((x, low_level_features), dim=)
         x = self.last_conv(x)
         x = F.interpolate(x, size=input.size()[:], mode='bilinear', align_corners=True)

         return x

     def _freeze_bn(self):
         for m in self.modules():
             if isinstance(m, BatchNorm2d):
                 m.eval()

     def _init_weight(self):
         for m in self.modules():
             if isinstance(m, nn.Conv2d):
                 n = m.kernel_size[] * m.kernel_size[] * m.out_channels
                 m.weight.data.normal_(, math.sqrt(. / n))
             elif isinstance(m, BatchNorm2d):
                 m.weight.data.fill_()
                 m.bias.data.zero_()

 def get_1x_lr_params(model):
     """
     This generator returns all the parameters of the net except for
     the last classification layer. Note that for each batchnorm layer,
     requires_grad is set to False in deeplab_resnet.py, therefore this function does not return
     any batchnorm parameter
     """
     b = [model.resnet_features]
     for i in range(len(b)):
         for k in b[i].parameters():
             if k.requires_grad:
                 yield k

 def get_10x_lr_params(model):
     """
     This generator returns all the parameters for the last layer of the net,
     which does the classification of pixel into classes
     """
     b = [model.aspp1, model.aspp2, model.aspp3, model.aspp4, model.conv1, model.conv2, model.last_conv]
     for j in range(len(b)):
         for k in b[j].parameters():
             if k.requires_grad:
                 yield k

 if __name__ == "__main__":
     model = DeepLabv3_plus(nInputChannels=, n_classes=, os=, pretrained=True, _print=True)
     model.eval()
     image = torch.randn(, , , )
     with torch.no_grad():
         output = model.forward(image)
     print(output.size())