Pytorch：注意力机制SELayer的代码实现

SE模块介绍

卷积层输出的特征图一般具有多个 plane，SENet 的主要思想是学习一个特征图中不同 plane 对于结果的影响权重，本质上就是在学习特征图不同 plane 间谁更重要，获得一组与特征图每个 plane 对应的权重系数向量，再利用该向量与原特征图向量矩阵进行点积，获得一个新的特征图，即 Squeeze-and-Excitation (SE) 模块。

Pytorch代码

import torch
import torch.nn as nn

# 直接在网络结构中调用就行
class SE(nn.Module):
    def __init__(self, inplanes, reduction=16):
        super(SE, self).__init__()
        self.se = nn.Sequential(
            nn.AdaptiveAvgPool2d(1), # 自适应平均池化，输出一维向量
            nn.Conv2d(inplanes, inplanes//reduction, 1, 1, 0),# 与下面注释效果一致
            # nn.Linear(inplanes, inplanes // reduction, bias=False),
            nn.ReLU(inplace=True),
            nn.Conv2d(inplanes // reduction, inplanes, 1, 1, 0),# 与下面注释效果一致
            # nn.Linear(inplanes // reduction, inplanes, bias=False),
            nn.Sigmoid()
        )


    def forward(self, x):
        input = x
        x = self.se(x)
        return input*x # 得到的权重向量与输入特征图相乘