FBSNet

Sun2024-10-072024-10-10

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

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

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

initial block
dual-branch backbone
feature aggregation module

Initial Block

Initial Block 包括三个的卷积层，并在每一个卷积层之后添加了 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)

Conv 1

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

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

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

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

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

Conv 2 and Conv 3

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

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

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

Batch Normalization

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

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

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

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

原理

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

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

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

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

注意事项

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

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

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

PReLU 激活函数

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

ReLU

ReLU 激活函数的定义是，其简单高效，计算速度快，但是存在以下问题:

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

PReLU

PReLU 定义如下:

$当当$

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

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

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

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)

Channel Attention Module (CAM)

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

其中，为通道注意力图，为输入特征，表示卷积核大小为的卷积运算

表示平均池化，表示压缩和重新加权，为 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)

首先通过一个自适应的2D平均池化层，对输入特征图在空间维度上进行全局平均池化，得到每个通道的全局平均值，将输入特征图的尺寸变为 ,每个通道通过池化被压缩成一个值，输出维度为 [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 是一个三分支模块，左分支负责提取局部信息和短距离信息，右分支用于扩大感受野以获取长距离特征信息，中间分支专门用于保存输入信息，示意图如下:

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

其代码如下：

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

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

左分支

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

右分支

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

中间分支

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

最后，输出重复开头的过程，进行批归一化和 PReLU 激活并经过卷积，再通过一个 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)

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

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

SDB

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

DRM

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

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

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

1
2
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)

SAM

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

其过程可描述如下：

其中，为所需的空间注意力图，为输入特征，表示卷积核大小为的卷积层，表示连接操作，表示平均池化操作，表示最大池化操作，为 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)

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

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

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

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

feature aggregation module

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

1	output = self.bn_prelu_8(output + output_sipath)

然后使用 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)

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

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

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

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

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

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