一、空间通道注意力机制
SpatialAttention 继承自 nn.Module,这是 PyTorch 中所有神经网络模块的基类。
__init__ 方法是构造函数,定义了各个层的参数。
kernel_size 是卷积核的尺寸,可以是 3 或 7。
padding 是根据卷积核尺寸计算的填充大小,保证输入输出的空间尺寸一致。
self.conv1 是一个卷积层,使用 1x1 的卷积核,将输入的通道数从 2 变为 1。
self.sigmoid 是一个 Sigmoid 激活函数,用于将输出值压缩到 [0, 1] 范围内,作为每个位置的注意力权重。
平均池化:
avg_out = torch.mean(x, dim=1, keepdim=True) 对输入特征图 x 在通道维度上进行平均池化,得到每个位置的平均值。keepdim=True 保持输出的维度与输入相同。
最大池化:
max_out, _ = torch.max(x, dim=1, keepdim=True) 对输入特征图 x 在通道维度上进行最大池化,得到每个位置的最大值。keepdim=True 保持输出的维度与输入相同。
特征拼接:
x = torch.cat([avg_out, max_out], dim=1) 将平均池化结果 avg_out 和最大池化结果 max_out 在通道维度上进行拼接,得到一个新的特征图 x,其通道数为 2。
卷积操作:
x = self.conv1(x) 使用 1x1 的卷积核对拼接后的特征图 x 进行卷积,将通道数从 2 变为 1。
注意力权重:
return self.sigmoid(x) 对卷积后的特征图 x 进行 Sigmoid 激活,得到每个位置的注意力权重。
二、提取掩码:
Translayer2_0, Translayer2_1, Translayer3_1, Translayer4_1:这些是不同分辨率特征图之间的转换层,将特征图的通道数转换为统一的 channel 数(例如 32)。
CFM:一个特征融合模块(Cross Feature Module),用于融合不同层次的特征。
ChannelAttention 和 SpatialAttention:分别是通道注意力模块和空间注意力模块,用于增强特征图的表示能力。
SAM:一个空间注意力模块(Self Attention Module),用于进一步处理特征图。
down05:一个下采样层,将特征图的分辨率缩小一半。
out_SAM 和 out_CFM:输出层,将特征图转换为单通道的预测图。