注意力机制——SEnet 代码实现及解释(Tensorflow and Pytorch)

Squeeze-and-Excitation Networks（挤压和激励网络）是一种通道注意力机制，类似一个小插件嵌入在深度学习网络中，通过对不同通道的权重重新标定来实现对重要信息的注意。PS：简单来说通过卷积得到的特征图就是所谓的通道，通道数量就是卷积层的过滤器数量。

提出SE模块的论文网址：https://openaccess.thecvf.com/content_cvpr_2018/html/Hu_Squeeze-and-Excitation_Networks_CVPR_2018_paper.html

代码地址：https://github.com/hujie-frank/SENet

SE模块原理解释：

我们通过结合下面的图像来理解SE的工作原理,以三维为例子进行解释，二维同样适用。心急的小伙伴可以直接划到下方使用代码哦。

1.在图的上方D×H×W×C是输入数据的形状，分别为深度×高度×宽度×通道数，二维则没有深度；

2.首先所有通道通过一个平均池化层将前三个维度均降维成1，输出形状为1×1×1×C。至此，挤压操作完成；

3.接下来输入到第一个全连接层中，该FC层采用的神经元为 $\frac{C}{ratio}$ ,ratio称为降维比率，由自己设置，一般为 $2^{n}$ ；

4.通过Relu函数计算第一个FC层的输出，即将负值全部置为0，正值保持不变;

5.连接上第二个FC层，该层使用的神经元数量为输入的通道数C，每一个输出值对应着一个原始通道；

6.最后通过sigmoid函数将所有输出值限定在（0,1）内，这些输出值对应的就是每一个通道的权重，最终通过与原始通道相乘(scale操作)实现权重的重标定。

SE代码实现：

#tensorflow代码实现,二维请根据注释修改一下即可使用
from tensorflow.keras.layers import GlobalAveragePooling3D, GlobalAveragePooling2D, Reshape, Dense, Multiply

#定义SE注意力机制
def squeeze_excite_block(input_tensor, ratio=16):
    # 挤压操作（全局平均池化）
    channel_axis = -1                             #这里获取的是通道数量,一般是最后一位数,如果实际中有不同请自行更改
    x = GlobalAveragePooling3D()(input_tensor)    #二维换成GlobalAveragePooling2D
    x = Reshape((1, 1, 1, x.shape[channel_axis]))(x)    #二维删掉一个(1,)
    # 激励操作（两个全连接层）
    x = Dense(x.shape[channel_axis] // ratio, activation='relu', kernel_initializer='he_normal', use_bias=True)(x)
    x = Dense(input_tensor.shape[channel_axis], activation='sigmoid', kernel_initializer='he_normal', use_bias=True)(x)
    # scale操作（将原始特征图与对应的权重相乘）
    x = Multiply()([input_tensor, x])
    return x

#使用SE模块
x = squeeze_excite_block(x, ratio=16)

#Pytorch代码实现SE,二维请根据注释修改一下即可使用
import torch
import torch.nn as nn
import torch.nn.functional as F

class squeeze_excite_block(nn.Module):
    def __init__(self, input_channels, ratio=16):
        super(squeeze_excite_block, self).__init__()
        #全局平均池化
        self.global_avg_pool = nn.AdaptiveAvgPool3d((1,1,1))   #二维换成nn.AdaptiveAvgPool2d((1,1))
        #挤压操作（两个FC层）
        self.fc1 = nn.Linear(input_channels, input_channels // ratio)
        self.relu = nn.ReLU(inplace=True)
        self.fc2 = nn.Linear(input_channels // ratio, input_channels)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        # 挤压操作（全局平均池化）
        b, c, d, h, w = x.size()                #二维换成b, c, h, w = x.size()
        y = self.global_avg_pool(x).view(b, c)
        # 激励操作（两个全连接层）
        y = self.relu(self.fc1(y))
        y = self.sigmoid(self.fc2(y))
        # 缩放操作（将输入张量与squeeze和excitation张量相乘）
        y = y.view(b, c, 1, 1, 1)               #二维换成y = y.view(b, c, 1, 1)
        return x * y

#使用SE模块
x = squeeze_excite_block(x)