一文看懂SE(Squeeze and Excitation)模块及代码实现

1,020 阅读2分钟

1. 简介

SE Block并非一个完整的网络,而且一个子结构,可以嵌在其他分类或检测模型中。SE Block的核心思想是通过网络根据loss去学习特征权重,使得有效的feature map权重更大,无效或效果小的feature map权重更小的方式去训练模型已达到更好的结果。它主要分为三部分:

Part1. Squeeze(压缩)操作:这一步骤通过全局平均池化(Global Average Pooling)将每个通道的空间信息(即高度和宽度)压缩成一个单独的值,从而获得一个包含所有通道全局信息的特征向量。这个向量的维度为1×1×C,其中C是通道数。

Part2. Excitation(激励)操作:在获得压缩后的特征向量后,通过两个全连接(Fully Connected,FC)层来预测每个通道的重要性权重。第一个FC层将特征向量的维度从C降低到C/r(r是降维比率),然后通过ReLU激活函数,第二个FC层再将维度从C/r增加回C。最后通过Sigmoid激活函数,得到每个通道的权重,这些权重反映了通道的重要性。

Part3. Scale(缩放)操作:将激励操作得到的通道权重应用到原始特征图上,每个通道的特征都乘以其对应的权重,从而对特征图进行重标定,增强了网络对有用特征的响应。

当然,SE block嵌在原有的一些分类网络中不可避免地增加了一些参数和计算量,但是在效果面前还是可以接受的。下面详细讲解其运算逻辑

2. 运算讲解

SE block示意图: SE示意图

Step1: 卷积操作(即图中的FtrF_{tr}操作)

严格来讲这一步是转换操作,并不是SE block的一部分,就是一个标准的卷积操作。输入输出定义如下:

FtrXU,XRWHC,URWHWF_{tr}: X \rightarrow U, X \in R^{W' * H' * C'}, U \in R^{W * H * W}

计算公式就是常规的卷积操作,计算公式如下:

uc=vcX=s=1Cvcsxsu_c = v_c * X = \sum ^{C'} _{s=1} v_c^{s} * x^s

其中,vcv_c 表示第c个卷积核, xsx^s 表示当前卷积核覆盖下的第s个输入, CC' 表示卷积核个数。

该操作得到了上图中左起第2个矩阵,其维度 = [H ,W, C]

Step2: FsqF_{sq}操作(即Squeenze操作)

该操作就是一个:global average pooling操作,公式如下:

zc=Fsq(uc)=1WHi=1Wj=1Huc(i,j)z_c = F_{sq}(u_c) = {1 \over W*H} \sum ^{W}_{i=1} \sum ^H_{j=1} u_c(i, j)

这里使用代码进行一定的解释: 代码如下:

x = torch.ones(size=(1, 2, 2, 3))
x[0][0][0][0] = 7
print("x = ", x)

avg_pool = torch.nn.AdaptiveAvgPool2d(1)    # 全局平均池化
x_pool = avg_pool(x)
print("x_pool.shape = ", x_pool.shape)
print("x_pool = ", x_pool)

输出结果: 在这里插入图片描述 计算的是每个通道的平均值,输出的shape=[1, 2, 1, 1]

这一步的结果相当于表明该层C个通道的数值分布情况,或者叫全局信息。

Step3: FexF_{ex}操作(即Excitation操作)

计算公式如下:

s=sigmoid(W2Relu(W1z))s = sigmoid(W_2 * Relu(W_1 z))

其中的zz表示上一步的zz, W1,W2W_1, W_2 表示的是线性层。这里计算出来的ss就是该模块的核心,用来表示各个channel的权重, 而且这个权重是通过前面这些全连接层和非线性层学习得到的,因此可以end-to-end训练。这两个全连接层的作用就是融合各通道的feature map信息,因为前面的squeeze都是在某个channel的feature map里面操作。

这里结合代码容易理解(即Pytorch实现SE模块):

class SELayer_2d(nn.Module):
    def __init__(self, channel, reduction=16):
        super(SELayer_2d, self).__init__()
        self.avg_pool = torch.nn.AdaptiveAvgPool2d(1)
        self.linear1 = nn.Sequential(
            nn.Linear(channel, channel // reduction, bias=False),
            nn.ReLU(inplace=True)
        )
        self.linear2 = nn.Sequential(
            nn.Linear(channel // reduction, channel, bias=False),
            nn.Sigmoid()
        )


    def forward(self, X_input):
        b, c, _, _ = X_input.size()  	# shape = [32, 64, 2000, 80]
        
        y = self.avg_pool(X_input)		# shape = [32, 64, 1, 1]
        y = y.view(b, c)				# shape = [32,64]
        
        # 第1个线性层(含激活函数),即公式中的W1,其维度是[channel, channer/16], 其中16是默认的
        y = self.linear1(y)				# shape = [32, 64] * [64, 4] = [32, 4]
        
        # 第2个线性层(含激活函数),即公式中的W2,其维度是[channel/16, channer], 其中16是默认的
        y = self.linear2(y) 			# shape = [32, 4] * [4, 64] = [32, 64]
        y = y.view(b, c, 1, 1)			# shape = [32, 64, 1, 1], 这个就表示上面公式的s, 即每个通道的权重

        return X_input*y.expand_as(X_input)

测试代码:

    data = torch.ones((32, 64, 2000, 80))
    se_2d = SELayer_2d(64)
    data_out = se_2d.forward(data)
    print("data_out = ", data_out.shape)

Step4: FscaleF_{scale}操作

计算公式如下:

x~=Fscale(uc,sc)=scuc\widetilde {x} = F_{scale}(u_c, s_c) = s_c ·u_c

其中,ucu_c表示uu中的一个通道,scs_c表示通道的权重。因此,相当于把每个通道的值乘以其权重。

代码即上述代码中的最后一行:

# y.expand_as(X_input)表示将y扩张到和X_input一样的维度
X_input*y.expand_as(X_input)		# 每个通道的值,乘以对应的权重

附录:

论文:Squeeze-and-Excitation Networks
论文链接:arxiv.org/abs/1709.01…
代码地址:github.com/hujie-frank…
PyTorch代码地址:github.com/miraclewkf/…

引用:

有参考添加链接描述,在此文章的理解上,增加了一些代码注释。