当前位置：首页 > article >正文

【深度学习】SAB：空间注意力

article 2025/1/15 0:06:42

@[toc]SAB：空间注意力

SAB：空间注意力

SA（Spatial Attention）空间注意力模块是一种用于计算机视觉任务的注意力机制，它的核心思想是通过学习特征图在空间维度上的重要性，自适应地增强或抑制不同空间位置的特征响应。与通道注意力（如 SENet 或 CBAM 中的通道注意力模块）不同，空间注意力关注的是特征图的空间维度（高度和宽度），而不是通道维度。

SA 的核心思想

1.空间注意力机制：

通过对特征图在通道维度上进行聚合（如最大池化或平均池化），生成空间注意力图。

使用卷积层学习空间位置之间的关系。

通过 Sigmoid 函数生成空间注意力权重，对原始特征图进行加权。

残差连接：

为了保留原始特征信息，通常会在空间注意力模块的输出上添加残差连接。

SA 的结构

通道聚合：

对输入特征图在通道维度上进行聚合（如最大池化或平均池化），生成一个空间描述图。

卷积层：

使用卷积层学习空间位置之间的关系。

通常使用一个 7x7 或 3x3 的卷积核来捕获局部空间信息。

Sigmoid 激活函数：

对卷积层的输出进行归一化，生成空间注意力权重。

特征加权：

将生成的空间注意力权重与输入特征图相乘，得到加权后的特征。

残差连接：

将加权后的特征与输入特征相加，保留原始信息。

代码实现

import torch
import torch.nn as nn

class SpatialAttention(nn.Module):
    def __init__(self, kernel_size=7):
        """
        初始化 SA 模块
        :param kernel_size: 卷积核大小，默认为 7
        """
        super(SpatialAttention, self).__init__()
        self.kernel_size = kernel_size

        # 通道聚合（最大池化和平均池化）
        self.max_pool = nn.AdaptiveMaxPool2d(1)
        self.avg_pool = nn.AdaptiveAvgPool2d(1)

        # 卷积层
        self.conv = nn.Conv2d(2, 1, kernel_size=kernel_size, padding=kernel_size // 2, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        """
        前向传播
        :param x: 输入特征图，形状为 [batch_size, channel, height, width]
        :return: 加权后的特征图
        """
        # 通道聚合
        max_out = self.max_pool(x)  # [batch_size, channel, 1, 1]
        avg_out = self.avg_pool(x)  # [batch_size, channel, 1, 1]

        # 拼接
        out = torch.cat([max_out, avg_out], dim=1)  # [batch_size, 2, height, width]

        # 卷积层
        out = self.conv(out)  # [batch_size, 1, height, width]

        # 生成空间注意力权重
        spatial_weights = self.sigmoid(out)  # [batch_size, 1, height, width]

        # 特征加权
        return x * spatial_weights