Squeeze-and-Excitation (SE) 模块是一种神经网络架构中的注意力机制,旨在通过学习每个通道的重要性来增强网络的表示能力。这种机制最早由 Jie Hu 等人在 2017 年的论文《Squeeze-and-Excitation Networks》中提出。SE 模块通过全局信息的“挤压”(squeeze)和每个通道的“激发”(excitation)来重新校准特征的通道权重。
SE 模块的工作原理
-
Squeeze(挤压):
- 对输入特征进行全局平均池化(Global Average Pooling),将每个通道的特征压缩成一个标量,捕获全局空间信息。
- 输出的维度为 C×1×1,其中 C 是通道数。
-
Excitation(激发):
- 利用两个全连接层和一个激活函数(ReLU 和 Sigmoid)来生成每个通道的权重。
- 通过这两个全连接层,SE 模块学习每个通道的重要性,输出的维度为 ( C )。
-
Recalibration(重新校准):
- 将生成的通道权重与原始特征逐通道相乘,重新校准特征图的通道权重。
PyTorch 实现
下面是一个简单的 PyTorch 实现 SE 模块的示例:
import torch
import torch.nn as nnclass SEModule(nn.Module):def __init__(self, channels, reduction=16):super(SEModule, self).__init__()self.avg_pool = nn.AdaptiveAvgPool2d(1)self.fc = nn.Sequential(nn.Linear(channels, channels // reduction, bias=False),nn.ReLU(inplace=True),nn.Linear(channels // reduction, channels, bias=False),nn.Sigmoid())def forward(self, x):b, c, _, _ = x.size()y = self.avg_pool(x).view(b, c)y = self.fc(y).view(b, c, 1, 1)return x * y.expand_as(x)# 示例使用
input_tensor = torch.randn(16, 64, 32, 32) # 批次大小为16,通道数为64,空间尺寸为32x32
se_module = SEModule(channels=64)
output_tensor = se_module(input_tensor)
print(output_tensor.shape) # 输出的形状应与输入相同
解释
-
SEModule
类:- 继承自
torch.nn.Module
。 - 在构造函数
__init__
中定义了全局平均池化层和两个全连接层。
- 继承自
-
forward
方法:- 对输入张量
x
进行全局平均池化,得到每个通道的全局信息。 - 将池化后的结果通过两个全连接层和激活函数,生成每个通道的权重。
- 将生成的权重与原始输入逐通道相乘,完成通道的重新校准。
- 对输入张量
-
示例使用:
- 创建一个输入张量,模拟批次大小为16,通道数为64,空间尺寸为32x32。
- 创建
SEModule
类的实例,指定通道数为64。 - 通过 SE 模块进行前向传播,得到重新校准后的输出张量。
通过这种方式,SE 模块能有效地提升网络的表示能力,特别是在图像分类、目标检测等任务中。