即插即用模块(3) -LSK 特征提取

原创于 2025-04-23 10:06:14 发布 · 1k 阅读

CC 4.0 BY-SA版权

文章标签：

在这里插入图片描述

paper：LSKNet: A Foundation Lightweight Backbone for Remote Sensing
Code：https://github.com/zcablii/LSKNet

大型选择性内核块 (LSK Block)

功能

通过动态调整感受野，自适应提取遥感图像中目标的上下文信息，增强目标与环境关联的建模能力，适用于外观相似但环境差异显著的场景。

实现流程
1. 大型内核卷积分解
  
  ：对输入 ( $X$ ) 应用深度可分离卷积，生成多尺度特征图序列 ( $U$ )。
2. 通道混合
  
  ：将 ( $U$ ) 拼接，通过 1x1 卷积生成特征图 ( $e$ )。
3. 空间内核选择
  
  ：对 ( e ) 进行平均池化和最大池化，生成描述符 ( A_{avg} ) 和 ( A_{max} )。
  
  拼接后通过卷积生成空间注意力图 ( SA )。
  
  应用 sigmoid 生成选择掩码 ( $SA_i$ )，加权融合生成注意力特征 ( S )。
4. 元素级乘法
  
  ：将 ( X ) 与 ( S ) 相乘，生成输出 ( Y )。
优势
- 动态调整感受野，适应不同目标的上下文需求。
- 深度可分离卷积降低计算复杂度。
- 增强目标与环境的特征关联，提升识别精度。

Large Selective Kernel Block结构图：

在这里插入图片描述

2、代码实现

import torch
import torch.nn as nn


class LSKblock(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.conv0 = nn.Conv2d(dim, dim, 5, padding=2, groups=dim)
        self.conv_spatial = nn.Conv2d(dim, dim, 7, stride=1, padding=9, groups=dim, dilation=3)
        self.conv1 = nn.Conv2d(dim, dim // 2, 1)
        self.conv2 = nn.Conv2d(dim, dim // 2, 1)
        self.conv_squeeze = nn.Conv2d(2, 2, 7, padding=3)
        self.conv = nn.Conv2d(dim // 2, dim, 1)

    def forward(self, x):
        attn1 = self.conv0(x)
        attn2 = self.conv_spatial(attn1)

        attn1 = self.conv1(attn1)
        attn2 = self.conv2(attn2)

        attn = torch.cat([attn1, attn2], dim=1)
        avg_attn = torch.mean(attn, dim=1, keepdim=True)
        max_attn, _ = torch.max(attn, dim=1, keepdim=True)
        agg = torch.cat([avg_attn, max_attn], dim=1)
        sig = self.conv_squeeze(agg).sigmoid()
        attn = attn1 * sig[:, 0, :, :].unsqueeze(1) + attn2 * sig[:, 1, :, :].unsqueeze(1)
        attn = self.conv(attn)
        return x * attn


if __name__ == '__main__':
    x = torch.randn(4, 64, 128, 128).cuda()
    model = LSKblock(64).cuda()
    out = model(x)
    print(out.shape)
).cuda()
    model = LSKblock(64).cuda()
    out = model(x)
    print(out.shape)