基于pytorch的DeepLabv3+语义分割实现

基于pytorch的DeepLabv3+语义分割实现

DeepLab系列从v1-v3+作为语义分割邻域中经典的网络模型，而V3+作为Deeplab所有思想的一个集合，实现Deeplabv3+也是入门语义分割邻域一个重要的知识点。

回顾DeepLabv1-v3

deepLabv1

DeepLabv1提出了深度卷积神经网络（DCNNs）与CRF相结合用于语义分割。DCNNs对目标做粗略分割，CRF做精细分割。由于本文主要针对深度卷积神经网络，因此只对深度卷积神经网络进行介绍。CRF作为DCNNs的后处理，在v3版本后的Deeplab系列也取消了CRF后处理机制。
DeepLabv1针对下采样会使得特征图尺寸减少，特征图尺寸减少对于语义分割效果有很大的影响。Deeplabv1通过修改最后几层的池化操作，在不改变特征图尺寸的前提下，提出空洞卷积，在不增加参数量的同时，获得更大的感受野，从而得到更加丰富的多尺度信息。

deepLabv2

DeepLabv2在v1的基础上针对分割目标具有多尺度大小的特征，使用不同空洞率的空洞卷积，并且借鉴SPP模型结构，提出空洞金字塔池化结构（ASPP)，解决对于多尺度目标分割的问题。
在DCNNs还是加入了CRF用于提高目标边界细节的分割。

deepLabv3

DeepLabv3的主要工作是改进了ASPP模块，提出串联或者并联ASPP模块中不同采样率的空洞卷积，并且在空洞卷积后加入BatchNormalize。并且文章中发现过大的采样率会使用33的空洞卷积退化为11卷积，因此将特征融合加入到ASPP模块中。

deepLabv3+

DeepLabv3+结合之前DeepLab系列的思想，提出编码-解码结构用于语义分割。在编码器阶段使用Xception结构以及ASPP结合作为图像的特征提取，在Xception模型中采用深度可分离卷积降低参数量。

DeepLabv3+中微调了传统Xception模型结构：
1、Xception模型分为入口模块、中间模块、出口模块三部分，v3+在Xception的中间模块进行了16次的堆叠操作，使得网络变得更深；
2、Xception模型中的所有池化操作被stride=2的深度可分离卷积所替代。
3、每个3*3深度可分离卷积后都加入BatchNormalize和ReLu。

基于pytorch的DeepLabv3+的网络模型实现

其中包括深度可分离卷积、ASPP、改进后的Xception模型中的入口模块、中间模块、出口模块。代码还是通过中文进行实现。训练和验证和之前文章一样，只需修改代码中的model即可。

import torch
import torch.nn as nn
import torch.nn.functional as F
import math


class 空洞金字塔池化模块ASPP(nn.Module):
    def __init__(self, 输入通道数, 输出通道数, 输出特征图相较于原始图像下降的比例):
        super(空洞金字塔池化模块ASPP, self).__init__()

        if 输出特征图相较于原始图像下降的比例 == 16:
            空洞率组合 = [1, 6, 12, 18]

        elif 输出特征图相较于原始图像下降的比例 == 8:
            空洞率组合 = [1, 12, 24, 36]

        self.ASPP中1x1卷积操作 = nn.Sequential(
            nn.Conv2d(输入通道数, 输出通道数, kernel_size=1, stride=1, padding=0, dilation=空洞率组合[0], bias=False),
            nn.BatchNorm2d(输出通道数),
            nn.ReLU()
        )

        self.ASPP中第一个3x3卷积操作 = nn.Sequential(
            nn.Conv2d(输入通道数, 输出通道数, kernel_size=3, stride=1, padding=空洞率组合[1], dilation=空洞率组合[1], bias=False),
            nn.BatchNorm2d(输出通道数),
            nn.ReLU()
        )

        self.ASPP中第二个3x3卷积操作 = nn.Sequential(
            nn.Conv2d(输入通道数, 输出通道数, kernel_size=3, stride=1, padding=空洞率组合[2], dilation=空洞率组合[2], bias=False),
            nn.BatchNorm2d(输出通道数),
            nn.ReLU()
        )

        self.ASPP中第三个3x3卷积操作 = nn.Sequential(
            nn.Conv2d(输入通道数, 输出通道数, kernel_size=3, stride=1, padding=空洞率组合[3], dilation=空洞率组合[3], bias=False),
            nn.BatchNorm2d(输出通道数),
            nn.ReLU()
        )

        self.全局池化操作 = nn.Sequential(
            nn.AdaptiveAvgPool2d((1, 1)),
            nn.Conv2d(输入通道数, 输出通道数, kernel_size=1, stride=1, padding=0, bias=False),
            nn.BatchNorm2d(输出通道数),
            nn.ReLU()
        )

        self.特征融合后1x1卷积操作 = nn.Sequential(
            nn.Conv2d(输出通道数*5, 256, 1, bias=False),
            nn.BatchNorm2d(256)
        )

        self._初始化权重()

    def forward(self, x):
        经过第一个ASPP中1x1卷积操作 = self.ASPP中1x1卷积操作(x)
        经过第一个ASPP中3x3卷积操作 = self.ASPP中第一个3x3卷积操作(x)
        经过第二个ASPP中3x3卷积操作 = self.ASPP中第二个3x3卷积操作(x)
        经过第三个ASPP中3x3卷积操作 = self.ASPP中第三个3x3卷积操作(x)
        经过ASPP中全局池化操作 = self.全局池化操作(x)
        经过ASPP中全局池化操作 = F.interpolate(经过ASPP中全局池化操作, size=经过第三个ASPP中3x3卷积操作.size()[2:], mode='bilinear', align_corners=True)

        拼接操作 = torch.cat((经过第一个ASPP中1x1卷积操作, 经过第一个ASPP中3x3卷积操作, 经过第二个ASPP中3x3卷积操作, 经过第三个ASPP中3x3卷积操作, 经过ASPP中全局池化操作), dim=1)
        输出 = self.特征融合后1x1卷积操作(拼接操作)
        return 输出

    def _初始化权重(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))

            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()


def 深度卷积中加入空洞卷积后的padding设置(输入图像, 卷积核尺寸, 空洞率):
    # 防止加入采样率后 卷积核的尺寸为偶数
    空洞率对应卷积核的尺寸 = 卷积核尺寸 + (卷积核尺寸 - 1) * (空洞率 - 1)
    padding_total = 空洞率对应卷积核的尺寸 - 1
    padding_begin = padding_total // 2
    padding_end = padding_total - padding_begin
    padding图像 = F.pad(输入图像, (padding_begin, padding_end, padding_begin, padding_end))

    return padding图像


class 深度可分离卷积(nn.Module):
    def __init__(self, 输入通道数, 输出通道数, 卷积核大小=3, 步长=1, 空洞率=1, bias=False):
        super(深度可分离卷积, self).__init__()

        self.深度卷积 = nn.Conv2d(输入通道数, 输入通道数, kernel_size=卷积核大小, stride=步长, padding=0, dilation=空洞率, groups=输入通道数, bias=bias)
        self.逐点卷积 = nn.Conv2d(输入通道数, 输出通道数, 1, 1, 0, 1, 1, bias=bias)

    def forward(self, x):
        padding = 深度卷积中加入空洞卷积后的padding设置(x, self.深度卷积.kernel_size[0], self.深度卷积.dilation[0])
        深度卷积操作 = self.深度卷积(padding)
        逐点卷积操作 = self.逐点卷积(深度卷积操作)
        return 逐点卷积操作


class Xception结构中重复使用的模块(nn.Module):
    def __init__(self, 输入通道数, 输出通道数, 每一个模块中相同操作的个数, 步长=1, 空洞率=1, 是否从RELU开始=True, 是否为出口模块=False):
        super(Xception结构中重复使用的模块, self).__init__()

        if 输出通道数 != 输入通道数 or 步长 != 1:
            self.跳跃连接 = nn.Sequential(
                nn.Conv2d(输入通道数, 输出通道数, 1, stride=步长, bias=False),
                nn.BatchNorm2d(输出通道数)
            )
        else:
            self.跳跃连接 = None

        self.relu = nn.ReLU(inplace=True)
        组成每一个模块 = []

        组成每一个模块.append(self.relu)
        组成每一个模块.append(深度可分离卷积(输入通道数, 输出通道数, 3, 步长=1, 空洞率=空洞率))
        组成每一个模块.append(nn.BatchNorm2d(输出通道数))
        对应每一次输出通道数的维度作为下一层输入通道数 = 输出通道数

        for 序号 in range(每一个模块中相同操作的个数 - 1):
            组成每一个模块.append(self.relu)
            组成每一个模块.append(深度可分离卷积(对应每一次输出通道数的维度作为下一层输入通道数, 对应每一次输出通道数的维度作为下一层输入通道数, 3, 步长=1, 空洞率=1))
            组成每一个模块.append(nn.BatchNorm2d(对应每一次输出通道数的维度作为下一层输入通道数))

        if not 是否从RELU开始:
            组成每一个模块[1:]

        if 步长 != 1:
            组成每一个模块.append(深度可分离卷积(输出通道数, 输出通道数, 3, 步长=2))

        if 步长 == 1 and 是否为出口模块:
            组成每一个模块.append(深度可分离卷积(输出通道数, 输出通道数, 3, 步长=1))

        self.组成每一个模块 = nn.Sequential(*组成每一个模块)

    def forward(self, input):
        x = self.组成每一个模块(input)

        if self.跳跃连接 is not None:
            跳跃连接 = self.跳跃连接(input)
        else:
            跳跃连接 = input

        x += 跳跃连接

        return x


class 出口模块(nn.Module):
    def __init__(self, 输入通道数, 输出通道数, 步长=1, 空洞率=1):
        super(出口模块, self).__init__()

        if 输出通道数 != 输入通道数 or 步长 != 1:
            self.跳跃连接 = nn.Sequential(
                nn.Conv2d(输入通道数, 输出通道数, 1, stride=步长, bias=False),
                nn.BatchNorm2d(输出通道数)
            )
        else:
            self.跳跃连接 = None

        self.relu = nn.ReLU(inplace=True)
        self.中间卷积层 = 深度可分离卷积(输入通道数, 输入通道数, 步长=步长, 空洞率=空洞率)
        self.最后输出层 = 深度可分离卷积(输入通道数, 输出通道数, 步长=步长, 空洞率=1)

    def forward(self, input):

        x = self.中间卷积层(input)
        x = self.relu(x)

        x = self.最后输出层(x)
        x = self.relu(x)

        if self.跳跃连接 is not None:
            跳跃连接 = self.跳跃连接(input)
        else:
            跳跃连接 = input

        x += 跳跃连接
        return x


class Xception(nn.Module):
    def __init__(self, 输入通道数=3, 输出特征图相较于原始图像下降的比例=16):
        super(Xception, self).__init__()

        if 输出特征图相较于原始图像下降的比例 == 16:
            入口模块中第三个模块的步长 = 2
            中间模块的空洞率 = 2
            出口模块的空洞率 = (1, 2)
        elif 输出特征图相较于原始图像下降的比例 == 8:
            入口模块中第三个模块的步长 = 1
            中间模块的空洞率 = 2
            出口模块的空洞率 = (2, 4)

        else:
            raise NotImplementedError

        self.入口模块中第一个卷积 = nn.Conv2d(输入通道数, 32, 3, stride=2, padding=1, bias=False)
        self.入口模块中第一个BN层 = nn.BatchNorm2d(32)
        self.relu = nn.ReLU(inplace=True)

        self.入口模块中第二个卷积 = nn.Conv2d(32, 64, 3, stride=1, padding=1, bias=False)
        self.入口模块中第二个BN层 = nn.BatchNorm2d(64)

        self.入口模块中第一个堆叠模块 = Xception结构中重复使用的模块(64, 128, 每一个模块中相同操作的个数=2, 步长=2, 是否从RELU开始=False)
        self.入口模块中第二个堆叠模块 = Xception结构中重复使用的模块(128, 256, 每一个模块中相同操作的个数=2, 步长=2, 是否从RELU开始=True)
        self.入口模块中第三个堆叠模块 = Xception结构中重复使用的模块(256, 728, 每一个模块中相同操作的个数=2, 步长=入口模块中第三个模块的步长, 是否从RELU开始=True)

        中间模块的第一个堆叠模块 = [Xception结构中重复使用的模块(728, 728, 每一个模块中相同操作的个数=3, 步长=1, 空洞率=中间模块的空洞率, 是否从RELU开始=True)] * 16
        self.中间模块的第一个堆叠模块 = nn.Sequential(*中间模块的第一个堆叠模块)

        self.出口模块 = 出口模块(728, 1024, 步长=1, 空洞率=出口模块的空洞率[1])

        self.出口卷积操作1 = nn.Sequential(
            深度可分离卷积(1024, 1536, 空洞率=出口模块的空洞率[1]),
            nn.BatchNorm2d(1536),
            nn.ReLU(inplace=True)
        )

        self.出口卷积操作2 = nn.Sequential(
            深度可分离卷积(1536, 1536, 空洞率=出口模块的空洞率[1]),
            nn.BatchNorm2d(1536),
            nn.ReLU(inplace=True)
        )

        self.出口卷积操作3 = nn.Sequential(
            深度可分离卷积(1536, 2048, 空洞率=出口模块的空洞率[1]),
            nn.BatchNorm2d(2048),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        x = self.入口模块中第一个卷积(x) # 1/2
        x = self.入口模块中第一个BN层(x)
        x = self.relu(x)

        x = self.入口模块中第二个卷积(x)
        x = self.入口模块中第二个BN层(x)
        x = self.relu(x)

        x = self.入口模块中第一个堆叠模块(x) # 1/4
        DCNN中直接取出的4倍下采样的特征图 = x
        x = self.入口模块中第二个堆叠模块(x)
        x = self.入口模块中第三个堆叠模块(x)

        x = self.中间模块的第一个堆叠模块(x)

        x = self.出口模块(x)
        x = self.出口卷积操作1(x)
        x = self.出口卷积操作2(x)
        x = self.出口卷积操作3(x)

        return x, DCNN中直接取出的4倍下采样的特征图


class DeepLabV3(nn.Module):
    def __init__(self, 输入图像的通道数=3, 分类个数=2, 输出特征图相较于输入特征图下降的比例=16):
        super(DeepLabV3, self).__init__()

        self.Xception提取特征 = Xception(输入图像的通道数, 输出特征图相较于原始图像下降的比例=输出特征图相较于输入特征图下降的比例)
        self.ASPP = 空洞金字塔池化模块ASPP(2048, 256, 输出特征图相较于原始图像下降的比例=16)

        self.出ASPP后的1x1卷积操作 = nn.Sequential(
            nn.Conv2d(256, 256, 1, bias=False),
            nn.BatchNorm2d(256),
            nn.ReLU())

        self.DCNN中下降4倍特征图的卷积操作 = nn.Sequential(
            nn.Conv2d(128, 48, 1, bias=False),
            nn.BatchNorm2d(48),
            nn.ReLU()
        )
        self.最后3x3卷积操作 = nn.Sequential(
            nn.Conv2d(304, 256, kernel_size=3, stride=1, padding=1, bias=False),
            nn.BatchNorm2d(256),
            nn.ReLU(),
            nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1, bias=False),
            nn.BatchNorm2d(256),
            nn.ReLU(),
            nn.Conv2d(256, 分类个数, kernel_size=1, stride=1)
        )

        self._初始化权重()

    def forward(self, input):
        x, DCNN4倍下采样的特征图 = self.Xception提取特征(input)
        x = self.ASPP(x)
        x = self.出ASPP后的1x1卷积操作(x)
        x = F.interpolate(x, size=(int(math.ceil(input.size()[-2] / 4)),
                                   int(math.ceil(input.size()[-1] / 4))), mode='bilinear', align_corners=True)

        DCNN4倍下采样的特征图 = self.DCNN中下降4倍特征图的卷积操作(DCNN4倍下采样的特征图)

        x = torch.cat((x, DCNN4倍下采样的特征图), dim=1)
        x = self.最后3x3卷积操作(x)
        x = F.interpolate(x, size=input.size()[2:], mode='bilinear', align_corners=True)

        return x

    def _初始化权重(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))

            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()