Squeeze-and-Excitation (SE) 模块

Squeeze-and-Excitation (SE) 模块是一种神经网络架构中的注意力机制，旨在通过学习每个通道的重要性来增强网络的表示能力。这种机制最早由 Jie Hu 等人在 2017 年的论文《Squeeze-and-Excitation Networks》中提出。SE 模块通过全局信息的“挤压”（squeeze）和每个通道的“激发”（excitation）来重新校准特征的通道权重。

SE 模块的工作原理

1. Squeeze（挤压）：

   • 对输入特征进行全局平均池化（Global Average Pooling），将每个通道的特征压缩成一个标量，捕获全局空间信息。
   • 输出的维度为 C×1×1，其中 C 是通道数。

2. Excitation（激发）：

   • 利用两个全连接层和一个激活函数（ReLU 和 Sigmoid）来生成每个通道的权重。
   • 通过这两个全连接层，SE 模块学习每个通道的重要性，输出的维度为 ( C )。

3. Recalibration（重新校准）：

   • 将生成的通道权重与原始特征逐通道相乘，重新校准特征图的通道权重。

PyTorch 实现

下面是一个简单的 PyTorch 实现 SE 模块的示例：

import torch
import torch.nn as nn

class 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)  # 输出的形状应与输入相同

解释

1. SEModule 类:

   • 继承自 torch.nn.Module。
   • 在构造函数 __init__ 中定义了全局平均池化层和两个全连接层。

2. forward 方法:

   • 对输入张量 x 进行全局平均池化，得到每个通道的全局信息。
   • 将池化后的结果通过两个全连接层和激活函数，生成每个通道的权重。
   • 将生成的权重与原始输入逐通道相乘，完成通道的重新校准。

3. 示例使用:

   • 创建一个输入张量，模拟批次大小为16，通道数为64，空间尺寸为32x32。
   • 创建 SEModule 类的实例，指定通道数为64。
   • 通过 SE 模块进行前向传播，得到重新校准后的输出张量。

通过这种方式，SE 模块能有效地提升网络的表示能力，特别是在图像分类、目标检测等任务中。