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

AI假装科研2025-04-27 15:10

paper：LSKNet: A Foundation Lightweight Backbone for Remote Sensing

Code：https://github.com/zcablii/LSKNet

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

功能

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

实现流程

大型内核卷积分解

：对输入 ( X X X ) 应用深度可分离卷积，生成多尺度特征图序列 ( U U U )。
通道混合

：将 ( U U U ) 拼接，通过 1x1 卷积生成特征图 ( e e e )。
空间内核选择

：对 ( e ) 进行平均池化和最大池化，生成描述符 ( A_{avg} ) 和 ( A_{max} )。

拼接后通过卷积生成空间注意力图 ( SA )。

应用 sigmoid 生成选择掩码 ( S A i SA_i SAi )，加权融合生成注意力特征 ( S )。
元素级乘法

：将 ( X ) 与 ( S ) 相乘，生成输出 ( Y )。

优势

动态调整感受野，适应不同目标的上下文需求。
深度可分离卷积降低计算复杂度。
增强目标与环境的特征关联，提升识别精度。

Large Selective Kernel Block结构图：

2、代码实现

python 复制代码

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)