2020CVPR-目标分割-SPNet-Strip Pooling Rethinking Spatial Pooling for Scene Parsing

论文下载：SPNet-paper:Strip Pooling: Rethinking Spatial Pooling for Scene Parsing
代码下载：SPNet-github:Strip Pooling: Rethinking Spatial Pooling for Scene Parsing

---

前言

Strip Pooling Rethinking Spatial Pooling for Scene Parsing是2020CVPR会议上比较有代表性的语义分割类文章，思路新型，具有较高的创新价值，无论是提出的SPM还是MPM模块，都可以即插即用，可以套用在各类基础网络上，后面会用代码解释。

一、创新点

在之前语义分割实验已经证明，空间池化可有效捕获用于场景分析等像素级预测任务的远程上下文信息。本文在传统的N×N型Spatial pooling的基础上，提出了一种考虑狭长卷积核（1×N或N×1）的strip pooling策略。并基于Strip pooling，进一步研究了Spatial pooling体系结构设计。通过对文章的总结，可以罗列以下创新点和贡献：

• 从新设计了池化结构，提出了非对称卷积1X N和N X 1的空间Pooling结构；
• 使用strip pooling设计了两个即插即用的Strip Pooling Module和Mixed Pooling Module；
• 基于strp pooling和mixed pooling搭建了spnet。

二、主要工作

1.背景

提高卷积神经网络中远程依赖关系建模能力的方法有：

• 引入注意力机制：self-attention机制或non-local模块。缺点：它们会消耗大量内存。
• 卷积角度上考虑：空洞卷积或者深度可分离卷积等，可以不引入额外参数同时扩大卷积神经网络的感受野。缺点：对小目标不友好
• 池化角度上考虑：全局/金字塔池化等，通过结合不同池化核大小的池化层，为图像提供全局信息。
  然而，空洞卷积和池化操作都是在正方形卷积中输入特征图并进行卷积运算。这限制了它们在捕获广泛存在于现实场景中的各向异性的上下文上的灵活性。
  例如，在某些情况下，目标对象可能具有长条形结构(下图b中的草地)或离散分布(下图a中的柱子)。使用大的方形池窗口不能很好地解决这个问题，因为它将不可避免地合并来自无关区域的污染信息。

为了更有效地捕获长依赖关系，本文在空间池化层扩大卷积神经网络感受野和捕获上下文信息的基础上，提出了条形池化（strip pooling）的概念。作为全局池化的替代方案，条纹池化有两个优点：

• 它沿着一个空间维度部署一个长条状的池化核形状，因此能够捕获孤立区域的长距离关系，如图1a和1c的第一行所示部分所示。
• 在其他空间维度上保持较窄的内核形状，便于捕获局部上下文，防止不相关区域干扰标签预测。

在网络中使用这种长而窄的池内核，可以使语义分割网络能够同时聚合全局和局部上下文信息。这是与传统的从固定的正方形区域收集上下文的池化有本质的不同。

三、核心内容

基于条纹池化的想法，作者提出了两种即插即用的池化模块 — Strip Pooling Module (SPM) 和 Mixed Pooling module (MPM)。

2.1、SPM——Strip Pooling Module

SPM由侧重于沿着水平和垂直空间两个维度捕获远程上下文的路径组成。

图中的条纹池化，实际上和普通池化方法没有区别，也是求池化核（长条形区域）所对应的特征图上位置的像素值求平均。
主要流程：

1. 输入一个特征图，这里实际上为C×H×W，为了方便讲解图中只画了一个通道。C个通道的特征图输入处理原理与这里所示的一个通道操作一模一样。
2. 输入的特征图经过水平和竖直条纹池化后变为H×1和1×W，对池化核内的元素值求平均，并以该值作为池化输出值。
3. 随后经过卷积核为3的1D卷积对两个输出feature map分别沿着左右和上下进行扩容，扩容后两个特征图尺寸相同，对扩容后的特征图对 应相同位置求和得到H×W的特征图。
4. 之后通过1×1的卷积与sigmoid处理后与原输入图对应像素相乘得到了输出结果。

在上面的过程中，输出张量中的每个位置都与输入张量中的各种位置建立了关系。例如，在上图中，输出张量中以黑框为界的正方形与所有与它具有相同水平或垂直坐标的位置相连(被红色和紫色边框包围)。因此，通过多次重复上述聚合过程，可以在整个场景中构建长期依赖关系。此外，得益于element-wise乘法操作，该SPM也可以被视为一种视觉注意力机制。

SPM可以直接应用于任何预先训练的骨干网络，而无需从无到有地进行训练。与全局平均池化相比，条纹池化考虑的是较长但较窄的范围，而不是整个特征图，避免了在相距较远的位置之间建立不必要的连接。与需要大量计算来建立每对位置之间关系的基于注意力的模块（no-local ）相比，SPM是轻量级的，可以很容易地嵌入到任何构建块中，从而提高捕获远程空间依赖关系和利用通道间依赖项的能力。
2.2代码的实现过程

由于原作者没有公开SPM代码，笔者根据文章描述写的代码，仅供参考，不对的地方欢迎大家一起讨论。

import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
import torch

class SPM(nn.Module):
    def __init__(self, input, in_channel, out_channel):
        super(SPM, self).__init__()
        inter_channel = in_channel / 4
        b, c, h, w = input.shape

        self.conv1 = nn.Sequential(F.conv2d(in_channel, inter_channel), nn.BatchNorm2d(inter_channel), nn.ReLU(True))

        self.h_pool = nn.AdaptiveAvgPool2d((1, None))
        self.w_pool = nn.AdaptiveAvgPool2d((None, 1))

        self.expand = F.upsample_bilinear((h, w))

        self.conv2 = nn.Sequential(F.conv2d(inter_channel, out_channel), nn.BatchNorm2d(out_channel), nn.ReLU(True))

        self.sigmoid = nn.Sigmoid()

    def forward(self, x):

        x_h = self.conv1(x)
        x_w = self.conv1(x)

        # split
        x_h = self.h_pool(x_h)
        x_w = self.w_pool(x_w)

        # expand
        x_h = self.expand(x_h)
        x_w = self.expand(x_w)

        # fusion
        fusion = x_h + x_w
        fusion = self.conv2(fusion)

        fusion = self.sigmoid(fusion)

        x = x.mul(fusion)

        return F.relu_(x)

if __name__=='main':
    image = Variable(torch.randn(4, 3, 256, 256))
    out = SPM(image, 64, 32)
    print(out)

**

2.2、MPM——Mixed Pooling Module

作者提出了一种新的附加组件块，称为混合池模块(MPM)，以进一步在高语义级别上建模长期依赖关系。它通过利用具有不同内核形状的池化操作来探测具有复杂场景的图像，从而收集有用的上下文信息。之前的研究结果表明，金字塔池模型(pyramid pooling module, PPM)是增强语义分割网络的有效方法。然而，PPM严重依赖于标准的池化操作（尽管不同的池内核位于不同的金字塔级别）。考虑到标准池化和条纹池化的优点，作者改进了PPM，提出了混合池模块(MPM)，它侧重于通过各种池化操作聚合不同类型的上下文信息，以使特征表示更有辨别力。优点是可以连续使用它来扩展远程依赖关系。在具有相同骨干网的情况下，仅具有两个MPM（原始PPM的大约1/3参数）的网络的性能甚至比PSPNet还要好。

提出的MPM由两个子模块组成，它们同时捕获不同位置之间的短距离和长距离依赖关系，作者发现这对于场景解析网络都是必不可少的。对于长距离依赖关系，与先前使用全局平均池化层的工作不同，文中通过同时使用水平和垂直条纹池化操作来捕获上下文信息。简化图可在上图（b）中找到。同时，条纹池化使在整个场景中离散分布的区域之间的连接和具有条纹状结构的编码区域成为可能。但是，对于语义区域分布紧密的情况，捕获局部上下文信息也需要进行空间池化。考虑到这一点，如上图（a）所示，文中采用了轻量级金字塔池子模块来进行短程依赖收集。它具有两个空间池化层，然后是用于多尺度特征提取的卷积层，以及用于原始空间信息保留的2D卷积层。每次合并后的特征图的大小分别为20×20和12×12，然后通过求和将所有三个子路径合并。基于以上两个子模块，文中将它们嵌套在具有瓶颈结构的残差块中，以进行参数缩减和模块化设计。具体地，在每个子模块之前，首先使用1×1卷积层来减少通道数量。然后将两个子模块的输出被串联在一起，并引入另一个1×1卷积层以进行通道扩展。其中，除了用于通道数量减少和扩展的卷积层以外，所有卷积层的内核大小均为3×3或3的倍数大小。

2.2 代码实现过程

同样这部分代码作者也没有公开，笔者只能根据文章描述写的代码，仅供参考，不对的地方欢迎大家一起讨论。

class MPM(nn.Module):

    def __init__(self, in_channels, pool_size, up_kwargs='bilinear', norm_layer=nn.BatchNorm2d):
        super(MPM, self).__init__()
        inter_channels = int(in_channels / 4)
        # 空间池化
        self.pool1 = nn.AdaptiveAvgPool2d(pool_size[0])
        self.pool2 = nn.AdaptiveAvgPool2d(pool_size[1])
        # strip pooling
        self.pool3 = nn.AdaptiveAvgPool2d((1, None))
        self.pool4 = nn.AdaptiveAvgPool2d((None, 1))

        self.conv1_1 = nn.Sequential(nn.Conv2d(in_channels, inter_channels, 1, bias=False),
                                     norm_layer(inter_channels),
                                     nonlinearity())
        self.conv1_2 = nn.Sequential(nn.Conv2d(in_channels, inter_channels, 1, bias=False),
                                     norm_layer(inter_channels),
                                     nonlinearity())

        self.conv2_0 = nn.Sequential(nn.Conv2d(inter_channels, inter_channels, 3, 1, 1, bias=False),
                                     norm_layer(inter_channels))

        self.conv2_1 = nn.Sequential(nn.Conv2d(inter_channels, inter_channels, 3, 1, 1, bias=False),
                                     norm_layer(inter_channels))
        self.conv2_2 = nn.Sequential(nn.Conv2d(inter_channels, inter_channels, 3, 1, 1, bias=False),
                                     norm_layer(inter_channels))

        self.conv2_3 = nn.Sequential(nn.Conv2d(inter_channels, inter_channels, (1, 3), 1, (0, 1), bias=False),
                                     norm_layer(inter_channels))
        self.conv2_4 = nn.Sequential(nn.Conv2d(inter_channels, inter_channels, (3, 1), 1, (1, 0), bias=False),
                                     norm_layer(inter_channels))

        self.conv2_5 = nn.Sequential(nn.Conv2d(inter_channels, inter_channels, 3, 1, 1, bias=False),
                                     norm_layer(inter_channels),
                                     nonlinearity())
        self.conv2_6 = nn.Sequential(nn.Conv2d(inter_channels, inter_channels, 3, 1, 1, bias=False),
                                     norm_layer(inter_channels),
                                     nonlinearity())
        self.conv3 = nn.Sequential(nn.Conv2d(inter_channels * 2, in_channels, 1, bias=False),
                                   norm_layer(in_channels))
        # bilinear interpolate options
        self._up_kwargs = up_kwargs

    def forward(self, x):
        b, c, h, w = x.size()
        x1 = self.conv1_1(x)
        x2 = self.conv1_2(x)
        x2_1 = self.conv2_0(x1)
        x2_2 = F.interpolate(self.conv2_1(self.pool1(x1)), size=[h, w], mode=self._up_kwargs)
        x2_3 = F.interpolate(self.conv2_2(self.pool2(x1)), size=[h, w], mode=self._up_kwargs)

        x2_4 = F.interpolate(self.conv2_3(self.pool3(x2)), size=[h, w], mode=self._up_kwargs)
        x2_5 = F.interpolate(self.conv2_4(self.pool4(x2)), size=[h, w], mode=self._up_kwargs)
        # PPM branch output
        x1 = self.conv2_5(F.relu_(x2_1 + x2_2 + x2_3))
        # MPM output
        x2 = self.conv2_6(F.relu_(x2_5 + x2_4))

        out = self.conv3(torch.cat([x1, x2], dim=1))
        return F.relu_(x + out)