Pytorch实现SE注意力机制_se注意力机制 pytorch

使用Pytorch实现SE通道注意力机制。

注意力机制就是在特征中添加权重，使得网络对重要特征特征进行关注，充分发挥重要特征的作用，从而提升网络模型的整体效果。
通道注意力机制SE就是在通道上添加权重，如果哪个通道更加重要，那么这个通道上的权重就会增加。

SE步骤

1. 对输入的特征图进行全局最大池化。输入特征图为(b,c,h,w),经过最大池化变成(b,c,1,1)，相当于把特征图变成了细长条。
2. 对经过全局平均池化的细长条进行全连接，先把通道缩小，比率为ratio。(b,c,1,1)->(b,c/ratio,1,1)
3. 再次进行全连接，使得通道数量恢复正常。(b,c/ratio,1,1)->(b,c,1,1)，这个就作为权重。
4. 将原始特征图与权重相乘输出。
   
   代码：

import torch
import torch.nn as nn

class SE(nn.Module):
    def __init__(self, channel,ratio = 4):
        super(SE,self).__init__()
        # 第一步：全局平均池化,输入维度(1,1,channel)
        self.avg = nn.AdaptiveAvgPool2d(1)
        # 第二步：全连接，通道数量缩减
        self.fc1 = nn.Linear(channel, channel//ratio,False)
        self.act = nn.ReLU()
        self.fc2 = nn.Linear(channel// ratio, channel,False)
        self.act2 = nn.Sigmoid()
    def forward(self,x):
        b, c, h, w = x.size()
        # b, c, 1, 1->b,c
        y = self.avg(x).view(b,c)
        y = self.fc1(y)
        y = self.act(y)
        y = self.fc2(y)
        y = self.act2(y).view(b,c,1,1)
        # print(y)
        return x * y

model = SE(8)
model = model.cuda()
input = torch.randn(1, 8, 12, 12).cuda()
output = model(input)
print(output.shape) # (1, 8, 12, 12)
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