（内含即插即用代码）选择性注意力LSK 模块，可用于改进YOLO系列的目标检测算法

最新推荐文章于 2025-05-11 13:54:05 发布

赢勾喜欢海

最新推荐文章于 2025-05-11 13:54:05 发布

阅读量1.8k

点赞数 15

文章标签： YOLO pytorch 目标检测计算机视觉人工智能

本文链接：https://blog.youkuaiyun.com/weixin_57111012/article/details/134448898

版权

文章详细讲解了如何利用LSKblock和Attention模块改进遥感目标检测算法，提供Python代码示例。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

有感兴趣的小伙伴可以搜一下这个相关的介绍资料，论文地址：https://arxiv.org/pdf/2303.09030.pdf

这一注意力可以有效用于遥感目标检测算法的改进，废话不多说，直接上代码;

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


class Attention(nn.Module):
    def __init__(self, d_model):
        super().__init__()

        self.proj_1 = nn.Conv2d(d_model, d_model, 1)
        self.activation = nn.GELU()
        self.spatial_gating_unit = LSKblock(d_model)
        self.proj_2 = nn.Conv2d(d_model, d_model, 1)

    def forward(self, x):
        shorcut = x.clone()
        x = self.proj_1(x)
        x = self.activation(x)
        x = self.spatial_gating_unit(x)
        x = self.proj_2(x)
        x = x + shorcut
        return x

在这里我们注意到，代码中是有LSKblock这个类和Attention这个类，无论哪一个，它的核心都是LSK的模块，无非是Attention加了一个残差结构。他们都是即插即用的模块，只需要一个形参，也就是输入通道数即可使用。

a = torch.ones(1,10,20,20)#设置输入
b = LSKblock(10)#实例化
c = Attention(10)#实例化
print(b(a).size())
print(c(a).size())

最后的结果如下，输入特征的尺寸并未改变