『深度学习笔记』CNN 网络设计中精巧通用的“小”插件

---

前言

• 所谓“插件”，就是要能锦上添花，又容易植入、落地，即真正的即插即用。本文盘点的“插件”能够提升 CNN 平移、旋转、scale 等变性能力或多尺度特征提取，感受野等能力，在很多 SOTA 网络中都会看到它们的影子。
• 所谓“插件”，就是不改变网络主体结构， 可以很容易嵌入到主流网络当中，提高网络提取特征的能力，能够做到plug-and-play

一、STN (空间变换器)

1.1 核心解析

• 在 OCR 等任务中，你会经常看到它的身影。对于 CNN 网络，我们希望其具有对物体的姿态、位置等有一定的不变性。即在测试集上可以适应一定的姿态、位置的变化。不变性或等变性可以有效提高模型泛化能力。虽然 CNN 使用 sliding-window 卷积操作，在一定程度上具有平移不变性。但很多研究发现，下采样会破坏网络的平移不变性。所以可以认为网络的不变性能力非常弱，更不用说旋转、尺度、光照等不变性。一般我们利用数据增强来实现网络的“不变性”。
• 本文提出 STN 模块，显式将空间变换植入到网络当中，进而提高网络的旋转、平移、尺度等不变性。可以理解为“对齐”操作。STN 的结构如上图所示，每一个 STN 模块由 Localisation net（本地网络），Grid generator（网格生成器） 和 Sampler（采样器） 三部分组成。

1. Localisation net 用于学习获取空间变换的参数，就是上式中的六个参数。
2. Grid generator 用于坐标映射。
3. Sampler 用于像素的采集，是利用双线性插值的方式进行。

1.2 核心代码

class SpatialTransformer(nn.Module):
    def __init__(self, spatial_dims):
        super(SpatialTransformer, self).__init__()
        self._h, self._w = spatial_dims 
        self.fc1 = nn.Linear(32*4*4, 1024) # 可根据自己的网络参数具体设置
        self.fc2 = nn.Linear(1024, 6)

    def forward(self, x): 
        batch_images = x #保存一份原始数据
        x = x.view(-1, 32*4*4)
        # 利用 FC 结构学习到 6 个参数
        x = self.fc1(x)
        x = self.fc2(x) 
        x = x.view(-1, 2,3) # 2x3
        # 利用 affine_grid 生成采样点
        affine_grid_points = F.affine_grid(x, torch.Size((x.size(0), self._in_ch, self._h, self._w)))
        # 将采样点作用到原始数据上
        rois = F.grid_sample(batch_images, affine_grid_points)
        return rois, affine_grid_points

二、ASPP

2.1 核心解析

• 本插件是带有空洞卷积的空间金字塔池化模块，主要是为了提高网络的感受野，并引入多尺度信息而提出的。我们知道，对于语义分割网络，通常面临是分辨率较大的图片，这就要求我们的网络有足够的感受野来覆盖到目标物体。对于 CNN 网络基本是靠卷积层的堆叠加上下采样操作来获取感受野的。本文的该模块可以在不改变特征图大小的同时控制感受野，这有利于提取多尺度信息。其中 rate 控制着感受野的大小，r 越大感受野越大。
• ASPP 主要包含以下几个部分：

1. 一个全局平均池化层得到 image-level 特征，并进行 1X1 卷积，并双线性插值到原始大小；
2. 一个 1X1 卷积层，以及三个 3X3 的空洞卷积；
3. 将 5 个不同尺度的特征在 channel 维度 concat 在一起，然后送入 1X1 的卷积进行融合输出。

2.2 核心代码

class ASPP(nn.Module):
    def __init__(self, in_channel=512, depth=256):
        super(ASPP,self).__init__()
        self.mean = nn.AdaptiveAvgPool2d((1, 1))
        self.conv = nn.Conv2d(in_channel, depth, 1, 1)
        self.atrous_block1 = nn.Conv2d(in_channel, depth, 1, 1)
        # 不同空洞率的卷积
        self.atrous_block6 = nn.Conv2d(in_channel, depth, 3, 1, padding=6, dilation=6)
        self.atrous_block12 = nn.Conv2d(in_channel, depth, 3, 1, padding=12, dilation=12)
        self.atrous_block18 = nn.Conv2d(in_channel, depth, 3, 1, padding=18, dilation=18)
        self.conv_1x1_output = nn.Conv2d(depth * 5, depth, 1, 1)

    def forward(self, x):
        size = x.shape[2:]
     	# 池化分支
        image_features = self.mean(x)
        image_features = self.conv(image_features)
        image_features = F.upsample(image_features, size=size, mode='bilinear')
     	# 不同空洞率的卷积
        atrous_block1 = self.atrous_block1(x)
        atrous_block6 = self.atrous_block6(x)
        atrous_block12 = self.atrous_block12(x)
        atrous_block18 = self.atrous_block18(x)
        # 汇合所有尺度的特征
     	x = torch.cat([image_features, atrous_block1, atrous_block6,atrous_block12, atrous_block18], dim=1)
        # 利用1X1卷积融合特征输出
        x = self.conv_1x1_output(x)
        return net

三、Non-local

3.1 核心解析

Non-Local 是一种 attention 机制，也是一个易于植入和集成的模块。Local 主要是针对感受野(receptive field)来说的，以 CNN 中的卷积操作和池化操作为例，它的感受野大小就是卷积核大小，而我们常用 3X3 的卷积层进行堆叠，它只考虑局部区域，都是 local 的运算。不同的是，non-local 操作感受野可以很大，可以是全局区域，而不是一个局部区域。捕获长距离依赖（long-range dependencies），即如何建立图像上两个有一定距离的像素之间的联系，是一种注意力机制。所谓注意力机制就是利用网络生成 saliency map，注意力对应的是显著性区域，是需要网络重点关注的区域。

具体实现步骤：

• 首先分别对输入的特征图进行 1X1的卷积来压缩通道数，得到 $\theta ,\varphi ,g$ 特征。
• 通过 reshape 操作，转化三个特征的维度，然后对进行矩阵乘操作，得到类似协方差矩阵, 这一步为了计算出特征中的自相关性，即得到每帧中每个像素对其他所有帧所有像素的关系。
• 然后对自相关特征进行 Softmax 操作，得到 0~1 的 weights，这里就是我们需要的 Self-attention 系数。
• 最后将 attention 系数，对应乘回特征矩阵 g 上，与原输入 feature map X 残差相加输出即可。

3.2 核心代码

class NonLocal(nn.Module):
    def __init__(self, channel):
        super(NonLocalBlock, self).__init__()
        self.inter_channel = channel // 2
        self.conv_phi = nn.Conv2d(channel, self.inter_channel, 1, 1,0, False)
        self.conv_theta = nn.Conv2d(channel, self.inter_channel, 1, 1,0, False)
        self.conv_g = nn.Conv2d(channel, self.inter_channel, 1, 1, 0, False)
        self.softmax = nn.Softmax(dim=1)
        self.conv_mask = nn.Conv2d(self.inter_channel, channel, 1, 1, 0, False)

    def forward(self, x):
        # [N, C, H , W]
        b, c, h, w = x.size()
        # 获取phi特征，维度为[N, C/2, H * W]，注意是要保留batch和通道维度的，是在HW上进行的
        x_phi = self.conv_phi(x).view(b, c, -1)
        # 获取theta特征，维度为[N, H * W, C/2]
        x_theta = self.conv_theta(x).view(b, c, -1).permute(0, 2, 1).contiguous()
        # 获取g特征，维度为[N, H * W, C/2]
        x_g = self.conv_g(x).view(b, c, -1).permute(0, 2, 1).contiguous()
        # 对phi和theta进行矩阵乘，[N, H * W, H * W]
        mul_theta_phi = torch.matmul(x_theta, x_phi)
        # softmax拉到0~1之间
        mul_theta_phi = self.softmax(mul_theta_phi)
        # 与g特征进行矩阵乘运算，[N, H * W, C/2]
        mul_theta_phi_g = torch.matmul(mul_theta_phi, x_g)
        # [N, C/2, H, W]
        mul_theta_phi_g = mul_theta_phi_g.permute(0, 2, 1).contiguous().view(b, self.inter_channel, h, w)
        # 1X1卷积扩充通道数
        mask = self.conv_mask(mul_theta_phi_g)
        out = mask + x # 残差连接
        return out

四、SE

4.1 核心解析

本文是 ImageNet 最后一届比赛的冠军作品，你会在很多经典网络结构中看到它的身影，例如 Mobilenet v3。其实是一种通道注意力机制。由于特征压缩和FC的存在，其捕获的通道注意力特征是具有全局信息的。本文提出了一种新的结构单元—— “Squeeze-and Excitation(SE)” 模块，可以自适应的调整各通道的特征响应值，对通道间的内部依赖关系进行建模。

有以下几个步骤：

• Squeeze: 沿着空间维度进行特征压缩，将每个二维的特征通道变成一个数，是具有全局的感受野。
• Excitation: 每个特征通道生成一个权重，用来代表该特征通道的重要程度。
• Reweight: 将 Excitation 输出的权重看做每个特征通道的重要性，通过相乘的方式作用于每一个通道上。

4.2 核心代码

class SE_Block(nn.Module):
    def __init__(self, ch_in, reduction=16):
        super(SE_Block, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)  # 全局自适应池化
        self.fc = nn.Sequential(
            nn.Linear(ch_in, ch_in // reduction, bias=False),
            nn.ReLU(inplace=True),
            nn.Linear(ch_in // reduction, ch_in, bias=False),
            nn.Sigmoid()
        )

    def forward(self, x):
        b, c, _, _ = x.size()
        y = self.avg_pool(x).view(b, c) # squeeze操作
        y = self.fc(y).view(b, c, 1, 1) # FC获取通道注意力权重，是具有全局信息的
        return x * y.expand_as(x) # 注意力作用每一个通道上

五、CBAM

5.1 核心解析

• SENet 在 feature map 的通道上进行 attention 权重获取，然后与原来的 feature map 相乘。这篇文章指出，该种 attention 方法只关注了通道层面上哪些层会具有更强的反馈能力，但是在空间维度上并不能体现出 attention。
• CBAM 作为本文的亮点，将 attention 同时运用在 channel 和 spatial 两个维度上， CBAM 与 SE Module 一样，可以嵌入在大部分的主流网络中，在不显著增加计算量和参数量的前提下能提升模型的特征提取能力。

通道注意力： 如上图所示

1. 首先，输入是一个 H×W×C 的特征 F，我们先分别进两个空间的全局平均池化和最大池化得到 两个 1×1×C 的通道描述。
2. 再将它们分别送进一个两层的神经网络，第一层神经元个数为 C/r，激活函数为 Relu，第二层神经元个数为 C。注意，这个两层的神经网络是共享的。
3. 然后，再将得到的两个特征相加后经过一个 Sigmoid 激活函数得到权重系数 Mc。
4. 最后，拿权重系数和 原来的特征 F 相乘即可得到缩放后的新特征。

伪代码:

def forward(self, x):
    # 利用FC获取全局信息，和Non-local的矩阵相乘本质上式一样的
    avg_out = self.fc2(self.relu1(self.fc1(self.avg_pool(x))))
    max_out = self.fc2(self.relu1(self.fc1(self.max_pool(x))))
    out = avg_out + max_out
    return self.sigmoid(out)

空间注意力：

• 与通道注意力相似，给定一个 H×W×C 的特征 F‘，我们先分别进行一个通道维度的平均池化和最大池化得到两个 H×W×1 的通道描述，并将这两个描述按照通道拼接在一起。然后，经过一个 7×7 的卷积层， 激活函数为 Sigmoid，得到权重系数 Ms。最后，拿权重系数和特征 F’ 相乘即可得到缩放后的新特征。

伪代码:

def forward(self, x):
    # 这里利用池化获取全局信息
    avg_out = torch.mean(x, dim=1, keepdim=True)
    max_out, _ = torch.max(x, dim=1, keepdim=True)
    x = torch.cat([avg_out, max_out], dim=1)
    x = self.conv1(x)
    return self.sigmoid(x)

待补全调整...

参考链接

1. https://mp.weixin.qq.com/s/glQSmbuvRI6CV8EyfYLezw