【python|attention】注意力机制代码

作者：Gausst松鼠会 | 2024-04-01 08:21:48

踩

注意力机制代码

every blog every motto: You can do more than you think.
https://blog.csdn.net/weixin_39190382?type=blog

0. 前言

梳理目前主流的注意力机制代码，目前以pytorch为例。
说明：

特征图维度的组织形式为：（batch,channel,height,width）
后续增加

1. 正文

1.1 SEBlock 2017

考虑通道间的注意力之间的关系，在通道上加入注意力机制
论文：https://arxiv.org/abs/1709.01507
代码：https://github.com/hujie-frank/SENet
在这里插入图片描述
对于输入特征图C2，其后加上SE注意力模块

1.1.1 步骤

主要分三步：

squeeze，对空间维度进行压缩，代码上即利用全局平均池化将每个通道平均成一个值，该值具有全局感受野。（通俗理解：一个通道理解为一个大饼，多个通道就是多个大饼垒在一起。全局平均池化即将一个大饼平均成一个点，整体看，类似一个垒起来的色子）
维度变化：（2,512,8,8）-> (2,512,1,1) ==> (2,512)
excitation, 利用权重学习上面各通道间的相关性，代码实现有全连接和卷积核为1的卷积操作两种方式。
维度变化：（2,512）-> (2,512//reducation)->(2,512) ==>(2,512,1,1)
说明： 该过程先降维在升维，降维倍数由reducation参数决定，降低网络的计算量，其中的激活函数增加了网络的非线性。
scale: 通过上面excitation的操作输出了每个通道的重要性，在通过乘法加权操作乘以输入数据C2，从而提升重要特征，抑制不重要特征。
维度变化：（2,512,8,8）*(2,512,1,1) -> (2,512,8,8)

小结： 即输入维度为（2,512,8,8），输出维度为：（2,512,8,8）
说明： 上述步骤中的“->”表示维度变化方向，“==>”表示通过view方法改变了维度。

1.1.2 更加清晰的理解图

在这里插入图片描述

说明：

全连接和1 × 1的卷积效果类似，上图显示为全连接，亦可为1*1的卷积，下同，不赘述。
激活函数位置见代码，下同。

1.1.3 代码

1. pytorch

class SELayer(nn.Module):
    def __init__(self, channel, reduction=4):
        """ SE注意力机制,输入x。输入输出特征图不变
            1.squeeze: 全局池化 (batch,channel,height,width) -> (batch,channel,1,1) ==> (batch,channel)
            2.excitaton: 全连接or卷积核为1的卷积(batch,channel)->(batch,channel//reduction)-> (batch,channel) ==> (batch,channel,1,1) 输出y
            3.scale: 完成对通道维度上原始特征的标定 y = x*y 输出维度和输入维度相同

        :param channel: 输入特征图的通道数
        :param reduction: 特征图通道的降低倍数
        """
        super(SELayer, self).__init__()
        # 自适应全局平均池化,即，每个通道进行平均池化，使输出特征图长宽为1
        self.avg_pool = nn.AdaptiveAvgPool2d(1)

        # 全连接的excitation
        self.fc = nn.Sequential(
            nn.Linear(channel, channel // reduction),
            nn.ReLU(inplace=True),
            nn.Linear(channel // reduction, channel),
            nn.Sigmoid()
        )
        # 卷积网络的excitation
        # 特征图变化：
        # (2,512,1,1) -> (2,512,1,1) -> (2,512,1,1)
        self.fc2 = nn.Sequential(
            nn.Conv2d(channel, channel // reduction, 1, bias=False),
            nn.ReLU(inplace=True),
            nn.Conv2d(channel // reduction, channel, 1, bias=False),
            nn.Sigmoid()
        )

    def forward(self, x):
        # (batch,channel,height,width) (2,512,8,8)
        b, c, _, _ = x.size()
        # 全局平均池化 (2,512,8,8) -> (2,512,1,1) -> (2,512)
        y = self.avg_pool(x).view(b, c)
        # (2,512) -> (2,512//reducation) -> (2,512) -> (2,512,1,1)
        y = self.fc(y).view(b, c, 1, 1)
        # (2,512,8,8)* (2,512,1,1) -> (2,512,8,8)
        pro = x * y
        return x * y
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41

2. tensorflow/keras

# SEBlock
feature_map_shape = input_x.shape  # input feature map shape
x = tf.reduce_mean(x, [1, 2])  # reduce along axis 1 and 2 ,height,width,
x = Dense(feature_map_shape[-1] / 16, activation=tf.nn.relu)(x)  # (batch,channel) -> (batch,channel/16)
x = Dense(feature_map_shape[-1], activation=tf.nn.relu)(x)  # (batch,channel/16) -> (batch,channel)
x = tf.multiply(input_x, x)  # multiply along channel
1
2
3
4
5
6

说明：

当使用全连接时，forward中平均池 (squeeze)，

# 全局平均池化 (2,512,8,8) -> (2,512,1,1) -> (2,512)
y = self.avg_pool(x).view(b, c)
1
2

当使用1*1卷积，forward中平均池化（squeeze）,

# 全局平均池化 (2,512,8,8) -> (2,512,1,1) 
y = self.avg_pool(x)
1
2

1.2 CBAM 2018

论文：https://arxiv.org/abs/1807.06521
代码：https://github.com/luuuyi/CBAM.PyTorch
CBAM可以无缝集成任何CNN架构中，开销不大，早期的注意力机制一种。
在这里插入图片描述
实验结果表明：顺序链接比并行连接好，其中通道注意力在前优于空间注意力在前。

1.2.1 通道注意力机制

1.2.1.1 概述

通道注意力机制和上面的SEBlock类似，唯一不同的是加了一个最大池化。而后，最大池化和平均池化共用一个多层感知机（mlp）, 再将结果相加和输入特征图进行点乘传入空间注意力机制。
说明： 主要步骤省略，可参考SEBlock和下面代码中的注释。
在这里插入图片描述

1.2.1.2 更加清晰的理解图

在这里插入图片描述

1.2.1.3 代码

说明： forward中，为了方便理解，展开书写了，等价于最开始注释了的几行。

class ChannelAttention(nn.Module):
    def __init__(self, in_channel, ratio=16):
        """ 通道注意力机制 同最大池化和平均池化两路分别提取信息，后共用一个多层感知机mlp,再将二者结合

        :param in_channel: 输入通道
        :param ratio: 通道降低倍率
        """
        super(ChannelAttention, self).__init__()
        # 平均池化
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        # 最大池化
        self.max_pool = nn.AdaptiveMaxPool2d(1)

        # 通道先降维后恢复到原来的维数
        self.fc1 = nn.Conv2d(in_channel, in_channel // ratio, 1, bias=False)
        self.relu1 = nn.ReLU()
        self.fc2 = nn.Conv2d(in_channel // ratio, in_channel, 1, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        # 平均池化
        # 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 x*self.sigmoid(out)
        
        # 平均池化一支 (2,512,8,8) -> (2,512,1,1) -> (2,512/ration,1,1) -> (2,512,1,1)
        # (2,512,8,8) -> (2,512,1,1)
        avg = self.avg_pool(x)
        # 多层感知机mlp (2,512,8,8) -> (2,512,1,1) -> (2,512/ration,1,1) -> (2,512,1,1)
        # (2,512,1,1) -> (2,512/ratio,1,1)
        avg = self.fc1(avg)
        avg = self.relu1(avg)
        # (2,512/ratio,1,1) -> (2,512,1,1)
        avg_out = self.fc2(avg)

        # 最大池化一支
        # (2,512,8,8) -> (2,512,1,1)
        max = self.max_pool(x)
        # 多层感知机
        # (2,512,1,1) -> (2,512/ratio,1,1)
        max = self.fc1(max)
        max = self.relu1(max)
        # (2,512/ratio,1,1) -> (2,512,1,1)
        max_out = self.fc2(max)

        # (2,512,1,1) + (2,512,1,1) -> (2,512,1,1)
        out = avg_out + max_out
        return x*self.sigmoid(out)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50

1.2.2 空间注意力机制

1.2.2.1 概述

将通道注意力机制的的结果作为新的输入特征图，主要步骤：

输入特征图，经过最大池化,平均池化（通道维度压缩，与前面的通道注意力机制不同）
维度变化：（2,512,8，8 ） -> (2,1,8,8)
将最大池化和平均池化在通道方向上合并
维度变化：（2,1,8,8）+ (2,1,8,8) -> (2,2,8,8)
3.经过卷积，通道变为1,再经过激活函数
维度变化：（2,2，8,8）-> (2,1,8,8)
和输入特征图点乘
维度变化：(2,512,8,8) * (2,1,8,8) -> (2,512,8,8)

1.2.2.2 更加清晰的理解图

说明： 下面两种图例其实是一个意思，即，单通道特征图，为了看起来更加清晰用了两种颜色。
在这里插入图片描述

1.2.2.3 代码

class SpatialAttention(nn.Module):
    def __init__(self, kernel_size=7):
        """ 空间注意力机制 将通道维度通过最大池化和平均池化进行压缩，然后合并，再经过卷积和激活函数，结果和输入特征图点乘

        :param kernel_size: 卷积核大小
        """
        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):
        print('x shape', x.shape)
        # (2,512,8,8) -> (2,1,8,8)
        avg_out = torch.mean(x, dim=1, keepdim=True)
        # (2,512,8,8) -> (2,1,8,8)
        max_out, _ = torch.max(x, dim=1, keepdim=True)
        # (2,1,8,8) + (2,1,8,8) -> (2,2,8,8)
        cat = torch.cat([avg_out, max_out], dim=1)
        # (2,2,8,8) -> (2,1,8,8)
        out = self.conv1(cat)
        return x * self.sigmoid(out)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26

声明：本文内容由网友自发贡献，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：【wpsshop博客】