FBSNet • Sun Blog

论文: https://arxiv.org/abs/2109.00699v1 ↗

代码: https://github.com/IVIPLab/FBSNet ↗

FBSNet 的网络结构可分为三部分:

initial block
dual-branch backbone
feature aggregation module

FBSNet

Initial Block#

Initial Block 包括三个 $3 \times 3$ 的卷积层，并在每一个卷积层之后添加了 Batch Normalization 和 PReLU 激活函数，在三层卷积层结束之后，又进行了一次 BN 和 PReLU

class Conv(nn.Module):
    def __init__(self, nIn, nOut, kSize, stride, padding, dilation=(1, 1), groups=1, bn_acti=False, bias=False):
        super().__init__()

        self.bn_acti = bn_acti

        self.conv = nn.Conv2d(nIn, nOut, kernel_size=kSize,
                              stride=stride, padding=padding,
                              dilation=dilation, groups=groups, bias=bias)

        if self.bn_acti:
            # Batch Normalization 和 PReLU 激活函数
            self.bn_prelu = BNPReLU(nOut)

    def forward(self, input):
        output = self.conv(input)

        if self.bn_acti:
            output = self.bn_prelu(output)

        return output


class BNPReLU(nn.Module):
    def __init__(self, nIn):
        super().__init__()
        self.bn = nn.BatchNorm2d(nIn, eps=1e-3)
        self.acti = nn.PReLU(nIn)

    def forward(self, input):
        output = self.bn(input)
        output = self.acti(output)

        return output

# --------------- class FBSNet ----------------------------
self.init_conv = nn.Sequential(
            Conv(3, 16, 3, 2, padding=1, bn_acti=True),
            Conv(16, 16, 3, 1, padding=1, bn_acti=True),
            Conv(16, 16, 3, 1, padding=1, bn_acti=True),
        )
        # 1/2
self.bn_prelu_1 = BNPReLU(16)

python

Conv 1#

Conv(3, 16, 3, 2, padding=1, bn_acti=True)

python

第一层的输入通道数为 3，输出通道数为 16，卷积核大小为 3，步长为 2，padding 为 1

由于步长为 2，这一层会对图像进行下采样，将输入的宽高减半

并通过 16 个 $3 \times 3$ 的卷积核对图像特征进行初步提取

在卷积操作后，进行批归一化和 PReLU 激活，提高模型的训练稳定性和非线性表达能力

Conv 2 and Conv 3#

Conv(16, 16, 3, 1, padding=1, bn_acti=True)
Conv(16, 16, 3, 1, padding=1, bn_acti=True)

python

第二层和第三层具有相同的结构，通过 16 个 $3 \times 3$ 的卷积核提取更高级的图像特征

但是这两层的步长均为 1，不会改变特征图的尺寸

Batch Normalization#

BN 是深度学习中一种用于加速神经网络训练，稳定网络训练过程的方法

在训练过程中，随着网络参数的更新，前一层的参数变化会导致后一层的输入分布发生变化，这种现象称为内部协变量偏移（Internal Covariate Shift）

内部协变量偏移会导致训练过程不稳定，收敛速度变慢，需要更小的学习率和更仔细的参数初始化

BN 通过对每个小批次的数据进行归一化，使得每一层的输入保持稳定的分布

原理#

首先计算批次的均值和方差

$\mu_B=\frac{1}{m}\sum_{i=1}^mx_i,\quad\sigma_B^2=\frac{1}{m}\sum_{i=1}^m(x_i-\mu_B)^2$

然后使用均值和方差对输入进行归一化，得到零均值，单位方差的输入

$\hat{x}_i=\frac{x_i-\mu_B}{\sqrt{\sigma_B^2+\epsilon}}$

这里的 $\epsilon$ 是一个很小的常数, 防止除以零

最后引入**可训练的参数 $\gamma$ 和 $\beta$ **对输入进行缩放和平移，恢复数据的表达能力

$y_i=\gamma\hat{x}_i+\beta$

注意事项#

BN 通常放在全连接层或卷积层之后，激活函数之前

对批大小敏感，批量太小可能导致估计的均值和方差不准确，影响模型性能

在训练阶段，使用当前批次的数据计算均值和方差，在测试阶段，使用在训练过程中累积的全局均值和方差

PReLU 激活函数#

PReLU（Parametric Rectified Linear Unit）是 ReLU（Rectified Linear Unit）激活函数的改进版本，它在 ReLU 的基础上增加了一个可学习的参数，用于调整负半轴的斜率。

ReLU#

ReLU 激活函数的定义是 $f(x) = max(0, x)$ ，其简单高效，计算速度快，但是存在以下问题:

死亡 ReLU 问题：当权重更新后，某些神经元可能永远输出 0，不再更新。
负半轴信息丢失：对于输入小于 0 的值，梯度为 0，可能导致有用的信息被忽略。

PReLU#

PReLU 定义如下:

$f(x)=\begin{cases}x,&\text{当 }x>0\\ax,&\text{当 }x\leq0\end{cases}$

其中 $a$ 是一个可学习的参数，初始值通常为一个小的正数

PReLU 由于负半轴的斜率 $a$ 可学习，且通常不为零，避免了神经元输出恒为零的情况

同时允许负值通过激活函数，保留了输入为负值时的有用信息

dual-branch backbone#

dual-branch backbone 采用双分支结构，由 Semantic Information Branch (SIB) 和 Spatial Detail Branch (SDB) 组成

SIB#

SIB 整体采用编码器—解码器结构，通过两个下采样块将特征图尺寸下采样到原始图像尺寸的 1/8，然后通过两个上采样块恢复原始图像尺寸

在下采样和上采样操作中间穿插了 Channel Attention Module (CAM) 和 Bottleneck Residual Unit (BRU)

SIB

Channel Attention Module (CAM)#

由于通道中包含丰富的特征信息和干扰噪声，SIB 使用 CAM 来强调需要突出的特征，在实现时采用 [ECA-Net](1910.03151 (arxiv.org) ↗) 中提出的 Efficient Channel Attention Module(ECA) 模块，其过程如下:

$M_{c}\left( X \right) =\sigma\left( C_{k\times k}\left( f_{Trans}\left( f_{AvgPool}(X) \right) \right) \right)$

其中， $M_c\in\mathbb{R}^{C\times1\times1}$ 为通道注意力图， $X\in\mathbb{R}^{C\times H \times W}$ 为输入特征， $C_{k \times k}$ 表示卷积核大小为 $k$ 的卷积运算

$f_{AugPool}(\cdot)$ 表示平均池化， $f_{Trans}(\cdot)$ 表示压缩和重新加权， $\sigma$ 为 Sigmoid 激活函数

其 pytorch 实现如下：

class eca_layer(nn.Module):
    """Constructs a ECA module.
    Args:
        channel: Number of channels of the input feature map
        k_size: Adaptive selection of kernel size
    """

    def __init__(self, channel, k_size=3):
        super(eca_layer, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.conv = nn.Conv1d(1, 1, kernel_size=k_size, padding=(k_size - 1) // 2, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        b, c, h, w = x.size()

        # feature descriptor on the global spatial information
        y = self.avg_pool(x)

        # Two different branches of ECA module
        y = self.conv(y.squeeze(-1).transpose(-1, -2)).transpose(-1, -2).unsqueeze(-1)

        # Multi-scale information fusion
        y = self.sigmoid(y)

        return x * y.expand_as(x)

python

首先通过一个自适应的2D平均池化层，对输入特征图在空间维度上进行全局平均池化，得到每个通道的全局平均值，将输入特征图的尺寸变为 $1 \times 1$ ,每个通道通过池化被压缩成一个值，输出维度为 [batch_size, channels, 1, 1]

然后通过 y.squeeze(-1).transpose(-1, -2) 去除最后一个维度并交换最后两个维度，得到 [batch_size, 1, channels]，从而将 1D 卷积操作应用到通道维度

然后再通过 transpose(-1, -2).unsqueeze(-1) 恢复原有形状 [batch_size, channels, 1, 1]

使用Sigmoid激活函数将输出值归一化到 0 到 1 之间，得到每个通道的权重

最后将得到的权重 y 与输入的特征图 x 相乘完成对通道的加权，调整各通道的贡献程度

Bottleneck Residual Unit (BRU)#

BRU 是一个三分支模块，左分支负责提取局部信息和短距离信息，右分支用于扩大感受野以获取长距离特征信息，中间分支专门用于保存输入信息，示意图如下:

BRU

图中使用的 $3 \times 1$ 和 $1 \times 3$ 卷积是采用卷积因子化策略，将标准的 $3 \times 3$ 卷积分解成两个一维卷积核，可以在保证模型性能的同时显著减少模型参数

其代码如下：

class BRUModule(nn.Module):
    def __init__(self, nIn, d=1, kSize=3, dkSize=3):  #
        super().__init__()
        self.bn_relu_1 = BNPReLU(nIn) 
        self.conv1x1_init = Conv(nIn, nIn // 2, 1, 1, padding=0, bn_acti=True)
        
        # -------------------------- 左分支 ----------------------------
        self.dconv3x1 = Conv(nIn // 2, nIn // 2, (dkSize, 1), 1, padding=(1, 0), groups=nIn // 2, bn_acti=True)
        self.dconv1x3 = Conv(nIn // 2, nIn // 2, (1, dkSize), 1, padding=(0, 1), groups=nIn // 2, bn_acti=True)
		self.ca11 = eca_layer(nIn // 2)
        self.dconv1x3_l = Conv(nIn // 2, nIn // 2, (1, dkSize), 1, padding=(0, 1), groups=nIn // 2, bn_acti=True)
        self.dconv3x1_l = Conv(nIn // 2, nIn // 2, (dkSize, 1), 1, padding=(1, 0), groups=nIn // 2, bn_acti=True)
        
		# -------------------------- 右分支 ----------------------------
        self.ddconv3x1 = Conv(nIn // 2, nIn // 2, (dkSize, 1), 1, padding=(1 * d, 0), dilation=(d, 1), groups=nIn // 2, bn_acti=True)
        self.ddconv1x3 = Conv(nIn // 2, nIn // 2, (1, dkSize), 1, padding=(0, 1 * d), dilation=(1, d), groups=nIn // 2, bn_acti=True)
        self.ca22 = eca_layer(nIn // 2)
        self.ddconv1x3_r = Conv(nIn // 2, nIn // 2, (1, dkSize), 1, padding=(0, 1 * d), dilation=(1, d), groups=nIn // 2, bn_acti=True)
        self.ddconv3x1_r = Conv(nIn // 2, nIn // 2, (dkSize, 1), 1, padding=(1 * d, 0), dilation=(d, 1), groups=nIn // 2, bn_acti=True)
        self.bn_relu_2 = BNPReLU(nIn // 2)
        
        # -------------------------- 中间分支 ----------------------------
        self.ca0 = eca_layer(nIn // 2)
        
        self.ca = eca_layer(nIn // 2)
        self.conv1x1 = Conv(nIn // 2, nIn, 1, 1, padding=0, bn_acti=False)
        self.shuffle_end = ShuffleBlock(groups=nIn // 2)

    def forward(self, input):
        output = self.bn_relu_1(input)
        output = self.conv1x1_init(output)

        br1 = self.dconv3x1(output)
        br1 = self.dconv1x3(br1)
        b1 = self.ca11(br1)
        br1 = self.dconv1x3_l(b1)
        br1 = self.dconv3x1_l(br1)

        br2 = self.ddconv3x1(output)
        br2 = self.ddconv1x3(br2)
        b2 = self.ca22(br2)
        br2 = self.ddconv1x3_r(b2)
        br2 = self.ddconv3x1_r(br2)


        output = br1 + br2 + self.ca0(output )+ b1 + b2

        output = self.bn_relu_2(output)

        output = self.conv1x1(output)
        output = self.ca(output)
        out = self.shuffle_end(output + input)
        return out

python

首先对输入进行批归一化和 PReLU 激活，经过 $1 \times 1$ 卷积减半通道数

左分支#

经过 (3,1) 和 (1,3) 的卷积操作，再经过 ECA 注意力层，并重复一次 (1,3) 和 (3,1) 卷积，提取局部信息和短距离信息

右分支#

使用膨胀卷积扩展感受野，再经过 ECA 注意力层并重复膨胀卷积，获取长距离特征信息

中间分支#

使用 ECA 层增强特征，然后再与左右分支 ECA 层的输出，最后的输出三者进行加和

最后，输出重复开头的过程，进行批归一化和 PReLU 激活并经过 $1 \times 1$ 卷积，再通过一个 ECA 层增强特征，然后通过 ShuffleBlock 进行通道混洗


class ShuffleBlock(nn.Module):
    def __init__(self, groups):
        super(ShuffleBlock, self).__init__()
        self.groups = groups

    def forward(self, x):
        '''Channel shuffle: [N,C,H,W] -> [N,g,C/g,H,W] -> [N,C/g,g,H,w] -> [N,C,H,W]'''
        N, C, H, W = x.size()
        g = self.groups
        #
        return x.view(N, g, int(C / g), H, W).permute(0, 2, 1, 3, 4).contiguous().view(N, C, H, W)

python

ShuffleBlock 将通道维度分为 g 组，每组的通道数为 c/g，然后使用 permute() 方法将通道和组的维度交换，最后恢复原始形状

从而将不同组的通道顺序打乱，使得下一层的卷积能够处理不同组之间混合后的通道，增加特征的表达能力和相互依赖性。

SDB#

在 SDB 分支中，使用了 Detail Residual Module (DRM) 和 Spatial Attention Module (SAM)

SDB

DRM#

DRM 是专门为补充语义分支中丢失的空间细节而设计的，由 3 个 $3 \times 3$ 的卷积层和一个 $1 \times 1$ 的卷积层构成

为了获得更多特征，将第二和第三个卷积层的通道数增加到原始输入的 4 倍（4C），最后使用一个 $1 \times 1$ 的卷积层将通道数再次减少到 C

DRM

但是在 FBSNet 开源的代码中，只使用了 3 个 $3 \times 3$ 的卷积层

self.conv_sipath1 = Conv(16, 32, 3, 1, 1, bn_acti=True)
self.conv_sipath2 = Conv(32, 128, 3, 1, 1, bn_acti=True)
self.conv_sipath3 = Conv(128, 32, 3, 1, 1, bn_acti=True)

python

SAM#

空间注意力模块 SAM 是沿通道轴应用最大池化和平均池化，然后通过标准卷积对其进行串联，从而生成有效的特征描述

其过程可描述如下：

$M_{s}\left( X \right) =\sigma\left( C_{k\times k}\left( \left[ f_{AvgPool}(X),f_{MaxPool}(X) \right] \right) \right)$

其中， $M_s\in\mathbb{R}^{1\times H\times W}$ 为所需的空间注意力图， $X\in\mathbb{R}^{C\times H\times W}$ 为输入特征， $C_{k \times k}$ 表示卷积核大小为 $k$ 的卷积层， $[]$ 表示连接操作， $f_{AugPool}(\cdot)$ 表示平均池化操作， $f_{MaxPool}(\cdot)$ 表示最大池化操作， $\sigma$ 为 Sigmoid 函数

其代码实现如下:

class SpatialAttention(nn.Module):
    def __init__(self, kernel_size=7):
        super(SpatialAttention, self).__init__()
        assert kernel_size in (3, 7), 'kernel size must be 3 or 7'
        padding = 3 if kernel_size == 7 else 1
        self.conv1 = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False)
        self.sigmoid = nn.Sigmoid()

    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)

python

首先沿着通道维度计算每个像素的通道均值，得到一个均值特征图，再沿着通道维度计算每个像素的通道最大值，得到一个最大值特征图

然后将均值特征图和最大值特征图沿着通道维度进行拼接，生成一个具有两个通道的特征图

对拼接后的特征图应用卷积操作，降维到单个通道，进一步提取空间特征

对卷积后的输出通过 Sigmoid 激活函数，将结果限制在 [0, 1] 之间，形成空间注意力权重矩阵

feature aggregation module#

首先将语义分支和空间分支的输出加和，并进行进行批归一化和 PReLU 激活

output = self.bn_prelu_8(output + output_sipath)

python

然后使用 CoordAttention ↗ 对特征进行融合，其代码如下

class CoordAtt(nn.Module):
    def __init__(self, inp, oup, reduction=4):
        super(CoordAtt, self).__init__()
        self.pool_h = nn.AdaptiveAvgPool2d((None, 1))
        self.pool_w = nn.AdaptiveAvgPool2d((1, None))

        mip = max(8, inp // reduction)

        self.conv1 = nn.Conv2d(inp, mip, kernel_size=1, stride=1, padding=0)
        self.bn1 = nn.BatchNorm2d(mip)
        self.act = h_swish()

        self.conv_h = nn.Conv2d(mip, oup, kernel_size=1, stride=1, padding=0)
        self.conv_w = nn.Conv2d(mip, oup, kernel_size=1, stride=1, padding=0)

    def forward(self, x):
        identity = x

        n, c, h, w = x.size()
        x_h = self.pool_h(x)
        x_w = self.pool_w(x).permute(0, 1, 3, 2)

        y = torch.cat([x_h, x_w], dim=2)
        y = self.conv1(y)
        y = self.bn1(y)
        y = self.act(y)

        x_h, x_w = torch.split(y, [h, w], dim=2)
        x_w = x_w.permute(0, 1, 3, 2)

        a_h = self.conv_h(x_h).sigmoid()
        a_w = self.conv_w(x_w).sigmoid()

        out = identity * a_w * a_h

        return out
    
class h_swish(nn.Module):
    def __init__(self, inplace=True):
        super(h_swish, self).__init__()
        self.sigmoid = h_sigmoid(inplace=inplace)

    def forward(self, x):
        return x * self.sigmoid(x)

python

首先对于输入特征，分别在 H 和 W 两个方向进行平均池化，得到两个大小分别为 [B, C, H, 1] 和 [B, C, W, 1](x_w 使用 permute() 做了转置)的特征图

然后将 x_h 和 x_w 沿着高度和宽度的方向拼接在一起，得到大小为 [B, C, H+W, 1] 的特征图

随后对其进行 $1 \times 1$ 卷积、批归一化和激活函数处理，融合特征信息

然后通过 torch.split() 将 y 在拼接的维度上分割，得到分别对应原始的高度和宽度方向的特征图

再分别通过 $1 \times 1$ 卷积得到高度和宽度方向的注意力权重，并通过 sigmoid() 激活函数将其压缩到 [0, 1] 范围，表示每个通道在高度和宽度方向的权重

最终的输出是输入乘以高度和宽度的权重，根据权重调整特征图的不同部分