语义分割网络DeepLab-v3的架构设计思想和TensorFlow实现

GitHub 地址:https://github.com/sthalles/deeplab_v3

语义分割

常规的图像分类深度卷积神经网络拥有相似的结构。这些模型以图像作为输入,并输出一个代表图像类别的数值。

通常,分类深度卷积神经网络有 4 种主要运算。卷积、激活函数、池化以及全连接层。传递一张图片,通过一系列这些运算会输出一个包含每个类别标签的概率的特征向量。请注意,在这种设定下,我们是对图片的整体进行分类。也就是说,为一张图像分配一个标签。

用于图像识别的标准深度学习模型。

与图像分类不同的是,在语义分割中我们将对图像中的每一个像素作出分类。所以,对每个像素而言,模型需要将它归类为预定义的类别之一。换言之,语义分割是在像素级别理解图像。

请记住,语义分割不会区分目标实例。因此,如果我们有同一个类的两个目标,它们最终将具有相同的类别标签。实例分割是区分同一类实例的问题。

语义分割与实例分割的区别。(中) 虽然它们是相同的目标,但它们被分类为不同的目标(实例分割)。(右) 相同的目标,相同的类别(语义分割)。

然而,常规的深度卷积神经网络 (如 AlexNet 和 VGG ) 并不适用于密集预测的任务。首先,这些模型包含许多用于减小输入特征的空间维度的层。结果,这些层最终产生缺乏清晰细节的高度抽象的特征向量。第二,全连接层在计算过程中具有固定的输入规模和松散的空间信息。

作为一个例子,试想通过一系列的卷积来传递图像,而不是使用池化和全连接层。我们将每次卷积都设置成步长为 1,padding 为「SAME」。通过这种处理,每一次卷积都保留了输入的空间维度。我们可以堆叠很多这种卷积,并最终得到一个分割模型。

用于密集预测的全卷积神经网络。请注意,不存在池化层和全连接层。

这个模型可以输出形状为 [W,H,C] 的概率张量,其中 W 和 H 代表的是宽度和高度,C 代表的是类别标签的个数。在第三个维度上应用最大化函数会得到形状为 [W,H,1] 的张量。然后,我们计算真实图像和我们的预测的每个像素之间的交叉熵。最终,我们对计算结果取平均值,并且使用反向传播算法训练网络。

然而,这个方法存在一个问题。正如前面所提到的,使用步长为 1,padding 为「SAME」,保留了输入的维度。但是,那样做的后果就是模型会极其耗费内存,而且计算复杂度也是很大的。

为了缓解这个问题,分割网络通常会有三个主要的组成部分:卷积层、降采样层和上采样层。

图像语义分割模型的编码器-解码器结构。

在卷积神经网络中实现降采样的常用方式有两个:通过改变卷积步长或者常规的池化操作。一般而言,降采样的目标就是减少给定特征图的空间维度。因此,降采样可以让我们在执行更深的卷积运算时不用过多地考虑内存。然而,这样一来在计算的时候会损失一些特征。

值得注意的是,这个架构的第一部分看上去类似于普通的分类深度卷积神经网络。不同的是,其中没有设置全连接层。

在第一部分之后,我们就得到了形状为 [W, H, D] 的特征向量,其中 W,H,D 分别是特征张量的宽度、高度和深度。注意,这个压缩向量的空间维度比原始输入更加少,但是更紧致。

顶部:VGG-16 网络的原始形式。要注意的是,堆叠的卷积层的顶部有三个全连接层。底部:VGG-16 网络中用 1x1 的卷积代替全连接层。这种改变可以让网络输出粗略的热图。

此时,常规的分类深度卷积神经网络会输出一个包含每个类别概率的密集(非空间)向量。取而代之,我们将这个压缩向量输入到一系列上采样层中。这些上采样层的作用就是重建与输入维度相同的输出向量。

通常,上采样都是基于步长置换卷积(strided transpose convolution)完成的。这些函数从深而窄的层变成浅而宽的层。这里,我们使用置换卷积将特征向量的维度增加到期望的结果。

在大多数论文中,分割网络的这两个部分被称作编码器和解码器。简而言之,第一部分将信息「编码」为压缩向量来代表输入。第二部分(解码器)的作用是将这个信号重建为期望的输出。

有很多基于编码器—解码器结构的神经网络实现。FCNs、SegNet,以及 UNet 是最流行的几个。

模型架构

与大多数编码器—解码器架构设计不同的是,Deeplab 提供了一种与众不同的语义分割方法。Deeplab 提出了一种用于控制信号抽取和学习多尺度语境特征的架构。

Deeplab 把在 ImagNet 上预训练得到的 ResNet 作为它的主要特征提取网络。但是,它为多尺度的特征学习添加了一个新的残差块。最后一个 ResNet 块使用了带洞卷积(atrous convolution),而不是常规的卷积。此外,这个残差块内的每个卷积都使用了不同的扩张率来捕捉多尺度的语境信息。

另外,这个残差块的顶部使用了带洞空间金字塔池化 (ASPP,Atrous Spatial Pyramid Pooling)。ASPP 使用了不同扩张率的卷积来对任意尺度的区域进行分类。

为了理解 Deeplab 的架构,我们需要着重注意这三个部分。(i)ResNet 架构,(ii) 带洞卷积,(iii) 带洞空间金字塔池化(ASPP)。接下来我们将逐一介绍这几个部分。

ResNets

ResNet 是一个非常流行的深度卷积神经网络架构,它赢得了 ILSVRC 2015 分类任务挑战赛的冠军。它的主要贡献之一就是提供了简化深度学习模型训练的框架。

在 ResNet 的原始形式中,它含有 4 个计算模块。每个模块包含不同数量的残差单元。这些单元以特别的形式执行一系列的卷积运算。同样,每个模块都夹杂了最大池化操作来减少空间维度。

原始论文提出了两种残差单元:基线块和瓶颈块。

基线块包含两个 3x3 的卷积,卷积中使用了 BN(批归一化)和 ReLU 激活函数。

残差模块。左:基线块;右:瓶颈块。

第二个是瓶颈块,它包括三个堆叠的部分,用一系列的 1x1、3x3 和 1x1 的卷积代替了之前的设计。两个 1x1 的卷积操作被用来减少和恢复维度。这使得中间的 3x3 的卷积可以在一个密度相对较低的特征向量上进行操作。此外,每个卷积之后、每个非线性 ReLU 之前都应用了 BN。

为了有助于澄清这个问题,我们将这一组操作定义为一个输入为 x 的函数 F——F(x)。

在 F(x) 中的非线性变换之后,这个单元将 F(x) 的结果和原始输入 x 相结合。这种结合是通过对两个函数求和得到的。原始输入 x 和非线性函数 F(x) 合并带来了一些优势。它使得前面的层可以访问后面层的梯度信号。换句话说,跳过 F(x) 上的操作允许前面的层访问更强的梯度信号。这种类型的连接已经被证明有助于更深网络的训练。

当我们增加模型容量时,非瓶颈单元也表明有助于准确率的提高。然而,瓶颈残差单元具有一些实际优势。首先,它们在几乎相同数量的参数下可执行更多的计算。第二,它们与非瓶颈单元的计算复杂度相似。

在实际中,瓶颈单元更适合于训练更深的模型,因为它们需要的训练时间和计算资源更少。

在我们的实现中,我们将使用完全预激活残差单元(full pre-activation Residual Unit),与标准瓶颈单元的唯一区别在于 BN 和 ReLU 激活函数的放置顺序。对于完全预激活,BN 和 ReLU(按此顺序)出现在卷积之前。

不同的 ResNet 构建模块块体系架构。最左边:原始 ResNet 模块;最右边:改进的完全预激活版本。

正如《Identity Mappings in Deep Residual Networks》中所展示的一样。完全预激活单元要优于其他的变体。

注意,这些设计之间的唯一区别是卷积堆栈中 BN 和 ReLU 的顺序。

带洞卷积

带洞卷积(或者扩张卷积)是具有一个因子的常规卷积,这个因子使得我们能够扩展滤波器的视野。

以 3×3 卷积滤波器为例。当扩张因子等于 1 时,它的行为类似于标准卷积。但是,如果将扩张因子设置为 2,则它具有扩大卷积核的效果。

理论上,它是这样工作的:首先,根据扩张率对卷积滤波器进行扩张。然后,它用零填充空白空间,创建稀疏的类似滤波器。最后,使用扩张的滤波器进行常规卷积。

不同扩张率的带洞卷积

因此,大小为 3x3、扩张率为 2 的卷积将使其能够覆盖 5x5 的区域。然而,因为它的作用就像一个稀疏的过滤器,只有原始的 3 x3 单元将执行计算并生成结果。

以类似的方式,扩张因子为 3 的常规 3×3 的卷积能够得到对应的 7×7 区域的信号。

这种效果允许我们控制计算特征响应的分辨率。此外,带洞卷积在不增加参数数量或计算量的情况下增加了更大范围的语境信息。

Deeplab 还表明,必须根据特征图的大小来调整扩张率。他们研究了在小特征图上使用大扩张率的结果。

给小特征图设置更大的扩张率的副作用。对于 14×14 的输入图像,使用扩张率为 15 的 3×3 卷积,其结果和常规的 1×1 卷积类似。

当扩张率非常接近特征图的尺寸时,一个常规的 3×3 的带洞滤波器的效果与标准的 1×1 卷积是一样的。

换句话说,带洞卷积的效率依赖于对扩张率的选择。由于这一原因,理解神经网络中的输出步长(output stride)的概念是很重要的。

输出步长反映输入图像大小与输出特征图大小的比率,它定义了输入向量在通过网络时经受的信号抽象程度。

输出步长为 16,图像大小为 224x224x3 时,输出特征向量比输入图像的维度小 16 倍,变成了 14x14。

此外,Deeplab 还讨论了不同输出步长对分割模型的影响。Deeplab 认为过强的信号抽象不利于密集预测任务。总之,具有较小输出步长 (较弱信号抽象) 的模型倾向于输出更精细的分割结果。然而,使用较小的输出步长训练模型需要更多的训练时间。

Deeplab 还展示了两种输出步长(8 和 16)设置下的结果。和预期的一样,步长等于 8 能够产生稍微好一些的结果。在这里,出于实际原因,我们选择了 16 为输出步长。

此外,由于带洞卷积块没有实现降采样,所以 ASPP 也运行在相同的特征响应大小上。因此,它允许使用相对较大的扩张率从多尺度的语境中学习特征。

新型带洞残差块包含三个残差单元。三个单元都总共拥有三个 3×3 的卷积块。在多重网格(multigrid)方法的启发下,Deeplab 为每个卷积设置了不同的扩张率。总之,多重网格为三个卷积中的每个卷积定义了不同的扩张率。

在实际中:

对于 block 4,当输出步长是 16,多重网格为(1,2,4)的时候,这三个卷积的扩张率分别是(2,4,8)。

带洞空间金字塔池化

带洞空间金字塔池化(ASPP)的思想是提供具有多尺度信息的模型。为了做到这一点,ASPP 添加了一系列具有不同扩张率的带洞卷积。这些扩张率是被设计用来捕捉大范围语境的。此外,为了增加全局的语境信息,ASPP 还通过全局平均池化(GAP)结合了图像级别的特征。

这个版本的 ASPP 包含 4 个并行的操作。它们分别是一个 1×1 的卷积以及三个 3×3 的卷积(扩张率分别是(6,12,18))。正如我们前面所提及的,现在,特征图的标称步长(nominal stride)是 16.

在原始实现的基础上,我们使用 513 x513 的裁剪尺寸进行训练和测试。因此,使用 16 的输出步长意味着 ASPP 接收大小为 32 x32 的特征向量。

此外,为了添加更多全局语境信息,ASPP 结合了图像级别的特征。首先,它将 GAP 应用于从最后一个带洞块输出的特征上。其次,所得特征被输入到具有 256 个滤波器的 1x 1 卷积中。最后,将结果进行双线性上采样到正确的维度大小。

  1. @slim.add_arg_scope

  2. def atrous_spatial_pyramid_pooling(net, scope, depth=256):

  3.    """

  4.    ASPP consists of (a) one 1×1 convolution and three 3×3 convolutions with rates = (6, 12, 18) when output stride = 16

  5.    (all with 256 filters and batch normalization), and (b) the image-level features as described in https://arxiv.org/abs/1706.05587

  6.    :param net: tensor of shape [BATCH_SIZE, WIDTH, HEIGHT, DEPTH]

  7.    :param scope: scope name of the aspp layer

  8.    :return: network layer with aspp applyed to it.

  9.    """

  10.    with tf.variable_scope(scope):

  11.        feature_map_size = tf.shape(net)

  12.        # apply global average pooling

  13.        image_level_features = tf.reduce_mean(net, [1, 2], name='image_level_global_pool', keep_dims=True)

  14.        image_level_features = slim.conv2d(image_level_features, depth, [1, 1], scope="image_level_conv_1x1", activation_fn=None)

  15.        image_level_features = tf.image.resize_bilinear(image_level_features, (feature_map_size[1], feature_map_size[2]))

  16.        at_pool1x1 = slim.conv2d(net, depth, [1, 1], scope="conv_1x1_0", activation_fn=None)

  17.        at_pool3x3_1 = slim.conv2d(net, depth, [3, 3], scope="conv_3x3_1", rate=6, activation_fn=None)

  18.        at_pool3x3_2 = slim.conv2d(net, depth, [3, 3], scope="conv_3x3_2", rate=12, activation_fn=None)

  19.        at_pool3x3_3 = slim.conv2d(net, depth, [3, 3], scope="conv_3x3_3", rate=18, activation_fn=None)

  20.        net = tf.concat((image_level_features, at_pool1x1, at_pool3x3_1, at_pool3x3_2, at_pool3x3_3), axis=3,

  21.                        name="concat")

  22.        net = slim.conv2d(net, depth, [1, 1], scope="conv_1x1_output", activation_fn=None)

  23.        return net

最后,各个分支的特征都被通过连接操作结合成一个单独的向量。然后使用另一个 1×1(采用 BN,和 256 个滤波器)的卷积对这个输出进行卷积。

ASPP 之后,我们将结果输入到另一个 1×1 的卷积中去生成最终的分割逻辑。

实现细节

这个实现用 ResNet-50 作为特征提取器,Deeplab_v3 采取了以下网络配置:

  • 输出步长=16

  • 为新的带洞残差块(block 4)使用固定的多重网格带洞卷积率(1,2,4)

  • 在最后一个带洞卷积残差块之后使用扩张率为(6,12,18)的 ASPP。

将输出步长设置为 16 有利于可持续地快速训练。与另一个输出步长 8 相比,输出步长为 16 使得带洞残差块处理的特征图比步长为 8 时处理的特征图小四倍。

将多重网格扩张率应用于带洞残差块内部的 3 个卷积。

最终,ASPP 中的三个并行的卷积得到了不同的扩张率——(6,12,18)。

在计算交叉熵损失函数之前,我们将分割逻辑调整为输入图像的大小。正如论文中所指出的,为了保持分辨率细节,调整分割逻辑的大小比调整真实标签的大小更好。

基于原始的训练过程,我们使用一个 0.5 到 2 之间的随机因子对每个图像做了扩展。此外,我们还对缩放后的图像做了随机的左右翻转。

最终,我们为训练和测试裁剪了 513 x513 大小的图像。

  1. def deeplab_v3(inputs, args, is_training, reuse):

  2.    # mean subtraction normalization

  3.    inputs = inputs - [_R_MEAN, _G_MEAN, _B_MEAN]

  4.    # inputs has shape [batch, 513, 513, 3]

  5.    with slim.arg_scope(resnet_utils.resnet_arg_scope(args.l2_regularizer, is_training,

  6.                                                      args.batch_norm_decay,

  7.                                                      args.batch_norm_epsilon)):

  8.        resnet = getattr(resnet_v2, args.resnet_model)

  9.        _, end_points = resnet(inputs,

  10.                               args.number_of_classes,

  11.                               is_training=is_training,

  12.                               global_pool=False,

  13.                               spatial_squeeze=False,

  14.                               output_stride=args.output_stride,

  15.                               reuse=reuse)

  16.        with tf.variable_scope("DeepLab_v3", reuse=reuse):

  17.            # get block 4 feature outputs

  18.            net = end_points[args.resnet_model + '/block4']

  19.            net = atrous_spatial_pyramid_pooling(net, "ASPP_layer", depth=256, reuse=reuse)

  20.            net = slim.conv2d(net, args.number_of_classes, [1, 1], activation_fn=None,

  21.                              normalizer_fn=None, scope='logits')

  22.            size = tf.shape(inputs)[1:3]

  23.            # resize the output logits to match the labels dimensions

  24.            # net = tf.image.resize_nearest_neighbor(net, size)

  25.            net = tf.image.resize_bilinear(net, size)

  26.            return net

为了实现残差网络 block4 中具有多重网格的带洞卷积,我们仅仅改变了 resnet_utils.py 文件中的以下这段代码:

  1. ...

  2. with tf.variable_scope('unit_%d' % (i + 1), values=[net]):

  3.  # If we have reached the target output_stride, then we need to employ

  4.  # atrous convolution with stride=1 and multiply the atrous rate by the

  5.  # current unit's stride for use in subsequent layers.

  6.  if output_stride is not None and current_stride == output_stride:

  7.    # Only uses atrous convolutions with multi-graid rates in the last (block4) block

  8.    if block.scope == "block4":

  9.      net = block.unit_fn(net, rate=rate * multi_grid[i], **dict(unit, stride=1))

  10.    else:

  11.      net = block.unit_fn(net, rate=rate, **dict(unit, stride=1))

  12.    rate *= unit.get('stride', 1)

训练

为了训练网络,我们决定使用来自于《Semantic contours from inverse detectors》的扩增版的 Pascal VOC 数据集。

训练数据由 8252 张图像组成。训练集有 5623 张,验证集有 2299 张。为了使用原始的 VOC2012 验证数据集来测试模型,我们从验证集中删去了 558 张图像。这 558 张图片也出现在官方的 VOC 验证集中。此外,我还添加了来自 VOC 2012 训练集中的 330 幅图像,它们既没出现在 5623 张训练集中,也没出现在 2299 张的验证集中。最后,8252 张图像中的 10%(大约 825 张图像)用来验证,其余的图像留着训练。

注意,这与原始论文是不一样的:这次实现没有在 COCO 数据集上预训练。此外,论文中描述到的一些训练和评估技术也没有用到。

结果

模型能够在 PASCAL VOC 验证集上得到良好的结果。

  • 像素准确率:大约 91%

  • 平均准确率:大约 82%

  • 均交并比(mIoU):大约 74%

  • 频权交并比(FWIoU):大约 86%

以下是 PASCAL VOC 验证集的图像分割的结果。

结论

语义分割无疑是计算机视觉领域中最流行的领域之一。Deeplab 提供了一个传统编码器-解码器体系架构的替代方案。它提倡在多范围的语境中使用带洞卷积学习特征。


原文链接:https://medium.freecodecamp.org/diving-into-deep-convolutional-semantic-segmentation-networks-and-deeplab-v3-4f094fa387df

理论NLPDeepLab-v3+语义分割实现TensorFlow
91