语义分割 | 轻量级实时分割经典BiSeNet及其进化

AI算法修炼营

发布于 2020-05-22 17:21:00

2.4K0

发布于 2020-05-22 17:21:00

轻量级语义分割

基于轻量化网络模型的设计作为一个热门的研究方法，许多研究者都在运算量、参数量和精度之间寻找平衡，希望使用尽量少的运算量和参数量的同时获得较高的模型精度。目前，轻量级模型主要有SqueezeNet、MobileNet系列和ShuffleNet系列等，这些模型在图像分类领域取得了不错的效果，可以作为基本的主干网络应用于语义分割任务当中。

然而，在语义分割领域，由于需要对输入图片进行逐像素的分类，运算量很大。通常，为了减少语义分割所产生的计算量，通常而言有两种方式：减小图片大小和降低模型复杂度。减小图片大小可以最直接地减少运算量，但是图像会丢失掉大量的细节从而影响精度。降低模型复杂度则会导致模型的特征提取能力减弱，从而影响分割精度。所以，如何在语义分割任务中应用轻量级模型，兼顾实时性和精度性能具有相当大的挑战性。

经典之作：BiseNet

论文地址：https://arxiv.org/abs/1808.00897.pdf

代码地址：https://github.com/CoinCheung/BiSeNet

本文对之前的实时性语义分割算法进行了总结，发现当前主要有三种加速方法：

1）通过剪裁或 resize 来限定输入大小，以降低计算复杂度。尽管这种方法简单而有效，空间细节的损失还是让预测打了折扣，尤其是边界部分，导致度量和可视化的精度下降；

2）通过减少网络通道数量加快处理速度，尤其是在骨干模型的早期阶段，但是这会弱化空间信息。

3）为追求极其紧凑的框架而丢弃模型的最后阶段（比如ENet）。该方法的缺点也很明显：由于 ENet 抛弃了最后阶段的下采样，模型的感受野不足以涵盖大物体，导致判别能力较差。

这些提速的方法会丢失很多 Spatial Details 或者牺牲 Spatial Capacity，从而导致精度大幅下降。为了弥补空间信息的丢失，有些算法会采用 U-shape 的方式恢复空间信息。但是，U-shape 会降低速度，同时很多丢失的信息并不能简单地通过融合浅层特征来恢复。

总结而言，实时性语义分割算法中，加速的同时也需要重视空间信息。论文中提出了一种新的双向分割网络BiSeNet。首先，设计了一个带有小步长的空间路径来保留空间位置信息生成高分辨率的特征图；同时设计了一个带有快速下采样率的语义路径来获取客观的感受野。在这两个模块之上引入一个新的特征融合模块将二者的特征图进行融合，实现速度和精度的平衡。

具体来说，空间路径Spatial Path使用较多的 Channel、较浅的网络来保留丰富的空间信息生成高分辨率特征；上下文路径Context Path使用较少的 Channel、较深的网络快速 downsample来获取充足的 Context。基于这两路网络的输出，文中还设计了一个Feature Fusion Module（FFM）来融合两种特征。

空间路径SP

减少下采样次数，只包含三个

的 Conv+BN+Relu，输出特征图的尺寸为原图的

。为了访存比考虑，此处并没有设计 Residual结构。由于它利用了较大尺度的特征图，所以可以编码比较丰富的空间信息。

class SpatialPath(nn.Module):
    def __init__(self, *args, **kwargs):
        super(SpatialPath, self).__init__()
        self.conv1 = ConvBNReLU(3, 64, ks=7, stride=2, padding=3)
        self.conv2 = ConvBNReLU(64, 64, ks=3, stride=2, padding=1)
        self.conv3 = ConvBNReLU(64, 64, ks=3, stride=2, padding=1)
        self.conv_out = ConvBNReLU(64, 128, ks=1, stride=1, padding=0)
        self.init_weight()

    def forward(self, x):
        feat = self.conv1(x)
        feat = self.conv2(feat)
        feat = self.conv3(feat)
        feat = self.conv_out(feat)
        return feat

上下文路径

上下文路径可以替换成任意的轻量网络，比如 Xception，ShuffleNet 系列，MobileNet 系列。本文主要采用 Xception39 和 ResNet-18 进行实验。

可以看到，为了准确率考虑，Context Path 这边使用了类似 U-shape 结构的设计，最终进行了32倍下采样。不过，不同于普通的 U-shape，此处只结合了最后两个 Stage，这样设计的原因主要是考虑速度。此外，和 DFN 类似，Context Path 依然在最后使用了 Global Average Pooling 来直接获取Global Context。

class ContextPath(nn.Module):
    def __init__(self, *args, **kwargs):
        super(ContextPath, self).__init__()
        self.resnet = Resnet18()
        self.arm16 = AttentionRefinementModule(256, 128)
        self.arm32 = AttentionRefinementModule(512, 128)
        self.conv_head32 = ConvBNReLU(128, 128, ks=3, stride=1, padding=1)
        self.conv_head16 = ConvBNReLU(128, 128, ks=3, stride=1, padding=1)
        self.conv_avg = ConvBNReLU(512, 128, ks=1, stride=1, padding=0)

        self.init_weight()

    def forward(self, x):
        H0, W0 = x.size()[2:]
        feat8, feat16, feat32 = self.resnet(x)
        H8, W8 = feat8.size()[2:]
        H16, W16 = feat16.size()[2:]
        H32, W32 = feat32.size()[2:]

        avg = F.avg_pool2d(feat32, feat32.size()[2:])
        avg = self.conv_avg(avg)
        avg_up = F.interpolate(avg, (H32, W32), mode='nearest')

        feat32_arm = self.arm32(feat32)
        feat32_sum = feat32_arm + avg_up
        feat32_up = F.interpolate(feat32_sum, (H16, W16), mode='nearest')
        feat32_up = self.conv_head32(feat32_up)

        feat16_arm = self.arm16(feat16)
        feat16_sum = feat16_arm + feat32_up
        feat16_up = F.interpolate(feat16_sum, (H8, W8), mode='nearest')
        feat16_up = self.conv_head16(feat16_up)

        return feat16_up, feat32_up # x8, x16

融合模块

在特征表示的层面上，两路网络的特征并不相同。因此不能简单地加权这些特征。由 Spatial Path 捕获的空间信息编码了绝大多数的丰富细节信息。而 Context Path 的输出特征主要编码语境信息。换言之，Spatial Path 的输出特征是低层级的，Context Path 的输出特征是高层级的。因此，提出一个独特的特征融合模块以融合这些特征。为了空间路径和上下文路径更好的融合，提出了特征融合模块FFM还有注意力优化模块ARM。

ARM:

ARM使用在上下文路径中，用于优化每一阶段的特征，使用全局平均池化指导特征学习，计算成本可以忽略。ARM应用全局平均池化来获取全局语义信息然后计算一个attention vector来知到特征学习。这个结构能够精细画Context Path中各个阶段的结果。它可以不用上采样就集成全局语义信息，计算代价较小。

class AttentionRefinementModule(nn.Module):
    def __init__(self, in_chan, out_chan, *args, **kwargs):
        super(AttentionRefinementModule, self).__init__()
        self.conv = ConvBNReLU(in_chan, out_chan, ks=3, stride=1, padding=1)
        self.conv_atten = nn.Conv2d(out_chan, out_chan, kernel_size= 1, bias=False)
        self.bn_atten = BatchNorm2d(out_chan, activation='none')
        self.sigmoid_atten = nn.Sigmoid()
        self.init_weight()

    def forward(self, x):
        feat = self.conv(x)
        atten = F.avg_pool2d(feat, feat.size()[2:])
        atten = self.conv_atten(atten)
        atten = self.bn_atten(atten)
        atten = self.sigmoid_atten(atten)
        out = torch.mul(feat, atten)
        return out

FFM:

在特征的不同层级给定的情况下，特征融合模块首先连接 Spatial Path 和 Context Path 的输出特征；接着，通过批归一化平衡特征的尺度。下一步，像 SENet 一样，把相连接的特征池化为一个特征向量，并计算一个权重向量。这一权重向量可以重新加权特征，起到特征选择和结合的作用。

将两个部分特征图通过concate方式叠加，然后使用类似SE模块的方式计算加权特征，起到特征选择和结合的作用。

class FeatureFusionModule(torch.nn.Module):
    def __init__(self, num_classes, in_channels):
        super().__init__()
        self.in_channels = in_channels
        self.convblock = ConvBlock(in_channels=self.in_channels,
                                   out_channels=num_classes,
                                   stride=1)
        self.conv1 = nn.Conv2d(num_classes, num_classes, kernel_size=1)
        self.relu = nn.ReLU()
        self.conv2 = nn.Conv2d(num_classes, num_classes, kernel_size=1)
        self.sigmoid = nn.Sigmoid()
        self.avgpool = nn.AdaptiveAvgPool2d(output_size=(1, 1))

    def forward(self, input_1, input_2):
        x = torch.cat((input_1, input_2), dim=1)
        assert self.in_channels == x.size(
            1), 'in_channels of ConvBlock should be {}'.format(x.size(1))
        feature = self.convblock(x)
        x = self.avgpool(feature)

        x = self.relu(self.conv1(x))
        x = self.sigmoid(self.conv2(x))
        x = torch.mul(feature, x)
        x = torch.add(x, feature)
        return x

损失函数：通过辅助损失函数监督模型的训练，通过主损失函数监督整个 BiSeNet 的输出。另外，还通过添加两个特殊的辅助损失函数监督 Context Path 的输出，就像多层监督一样。上述所有损失函数都是 Softmax。最后借助参数 α 以平衡主损失函数与辅助损失函数的权重。

可以看到，BiSeNet是一种很有效的设计。当替换上大模型之后，精度甚至高于 PSPNet 等算法。BiSeNet 算法对实时性语义分割算法提出了新的思考，在提升速度的同时也需要关注空间信息。同时，该设计也是一次对 Segmentation Backbone 的思考，希望设计一个对 Segmentation 任务友好的框架，当然现在还存在许多需要改进的地方。此外，该方法不仅仅可应用于实时性语义分割算法，也可应用于其他领域，尤其是在对 Spatial Detail 和 Context 同时有需求的情况下。并已有研究将其应用于 Potrait Segmentation。

BiSeNet升级版——BiSeNet V2

论文地址：

https://arxiv.org/abs/2004.02147

代码地址：

https://github.com/MaybeShewill-CV/bisenetv2-tensorflow

与几种最新的实时语义分割方法相比，BiSeNet V2具有良好的性能。具体来说，对于2048x1,024的输入，BiseNet2在Cityscapes测试集中的平均IoU达到72.6％，在一张NVIDIA GeForce GTX 1080 Ti卡上的速度为156 FPS，这比现有方法要快得多，而且可以实现更好的分割精度。

低层次的细节和高层次的语义是语义分割的基础。然而，为了加快模型推理的速度，目前的方法几乎总是牺牲低层次的细节，这导致了相当大的精度下降。

BiSeNet V2将这些空间细节和分类语义分开处理，以实现高精度和高效率的实时语义分割。为此，提出了一个有效的架构，在速度和精度之间进行权衡，称为双边分割网络(BiSeNet V2)。该体系结构包括:(1)一个细节分支，具有宽通道和浅层，用于捕获低层细节并生成高分辨率的特征表示;(2)一个语义分支，通道窄，层次深，获取高层次语义语境。语义分支是轻量级的，因为它减少了通道容量和快速下采样策略。此外，设计了一个引导聚合层来增强相互连接和融合这两种类型的特征表示。此外，还设计了一种增强型训练策略，在不增加任何推理代价的情况下提高分割性能。

双边分割网络概况主要有三个组成部分：紫色虚线框内的双通道主干，橙色虚线框内的聚集层，黄色虚线框内的助推部分。双通道主干有一个细节分支(蓝色的数据集)和一个语义分支(绿色的数据集)。三个阶段支路分别有C1、C2、C3通道。相应阶段的渠道语义分支可以轻量级的因子λ(λ< 1)。语义分支的最后一个阶段是上下文嵌入块的输出。同时，立方体中的数字是特征映射大小与输入分辨率的比值。在聚合层部分，我们采用了双边聚合层。将采样操作表明,Up代表upsampling操作,ϕ是Sigmoid函数,和x意味着element-wise输出。此外，在推理部分，设计了一些辅助分割头，以提高分割性能，没有任何额外的推理成本。

表1、细节分支和语义分支的例示。每个阶段S包含一个或多个操作opr(例如，Conv2d, Stem, GE, CE)。每个操作有一个大小为k的内核，步长为s，输出通道为c，重复r次。扩展因子e用于扩展操作的通道数。这里的通道比λ= 1/4。在细节分支的对应阶段，绿色标志着语义分支的通道更少。注:Conv2d表示卷积层，后面是一个批处理的归一化层和relu激活函数。阀杆表示阀杆块。GE代表采集-扩展层。CE是上下文嵌入块。

1、细节分支

细节分支负责空间细节，这是低级的信息。因此，该分支需要丰富的信道容量来编码丰富的空间细节信息。同时，因为细节分支只关注底层细节，所以我们可以为这个分支设计一个小跨度的浅层结构。总体而言，细节分支的关键概念是使用宽通道和浅层来处理空间细节。此外，该分支的特征表示具有较大的空间尺寸和较宽的信道。因此，最好不要采用残差连接，这样会增加内存访问成本，降低速度。

表1中细节分支的实例化包含三个阶段，每一层都是卷积层，然后是batch normalization和激活函数。每个阶段的第一层有一个stride s = 2，而同一阶段的其他层有相同数量的卷积和输出feature map大小。因此，这个分支提取的输出特征映射是原始输入的1/8。由于通道容量大，这个细节分支编码了丰富的空间细节。同时，由于通道容量大，空间维度大，残差结构将增加内存访问成本。因此，这个分支主要遵循的VGG网络原理来堆叠层。

2、语义分支

与细节分支并行，语义分支旨在捕获高级语义。该分支的信道容量较低，而空间细节可以由细节分支提供。相反,在我们的实验中,语义分支占比λ(λ<1)细节分支的通道，这使得这个分支是轻量级的。实际上，语义分支可以是任何轻量级的卷积模型。同时，语义分支采用快速下采样策略提高了特征表示的层次，快速扩大了接受域。高级语义需要大量的接受域。因此，语义分支使用全局平均池嵌入全局上下文响应。

考虑到接受域大，计算效率高，设计语义分支，其灵感来自轻量级图像分类模型的理念，如Xception、MobileNet、ShuffleNet 。语义分支的一些关键特性如下：

Stem Block：采用Stem Block作为语义分支的第一阶段，如图（a）所示。它使用两种不同的下采样方式来缩小特征表示。然后将两个分支的输出特性串联起来作为输出。该结构具有高效的计算成本和有效的特征表达能力。

Context Embedding Block：语义分支需要大的接受域来捕获高级语义。所以设计了Context Embedding Block。该块使用全局平均池和残差连接有效地嵌入全局上下文信息，如图（b）所示。

Gather-and-Expansion Layer：利用深度卷积的优点，提出了Gather-and-Expansion Layer，如图所示。主要包括:

(1)一个3×3的卷积，有效地聚合特征响应并扩展到高维空间;

(2)在膨胀层的每个单独输出通道上独立进行3×3深度卷积;

(3)以1×1的卷积作为投影层，将深度卷积的输出投影到低信道容量空间中。当stide = 2时，我们采用两个3×3的深度卷积，进一步扩大了感受野，一个3×3的可分离卷积用于shortcut。在这一层，用两个3×3深度卷积代替可分离变量卷积中的5×5深度卷积，这两个3×3深度卷积的计算量少，感受野相同。

与MobileNetv2的反向瓶颈层相比，GE层多了一个3×3的卷积，却有利于计算成本和内存访问成本，因为在CUDNN库中对3×3卷积进行了特别优化同时，由于这一层，GE层比反向瓶颈层具有更高的特征表达能力。

3、聚合层

细节分支和语义分支的特征表示是互补的，其中一个不知道另一个的信息。因此，设计了一个聚合层来合并这两种类型的特性表示。由于采用了快速下采样策略，语义分支的输出空间维度比细节分支小，所以需要对语义的输出特征图进行向上采样分支以匹配细节分支的输出。

有一些不同的方式来合并两种类型的特征响应，即元素方式的求和和连接。但是，这两个分支的输出具有不同级别的特征表示。细节分支是低级的，而语义分支是高级的。因此，简单的组合忽略了这两类信息的多样性，导致性能下降和难以优化。

在观察的基础上，我们提出双边引导的聚合层融合来自两个分支的互补信息，如图所示。该层使用的是上下文信息语义分支用来指导细节分支的特征响应。通过不同的尺度指导，可以捕获不同的尺度特征表示，这些特征表示对多尺度信息进行了固有的编码。同时，与简单的组合方式相比，这种引导方式可以使两个分支之间进行有效的通信。