【Attention】CBAM中的SAM

最新推荐文章于 2024-01-02 17:02:25 发布

essenge

最新推荐文章于 2024-01-02 17:02:25 发布

阅读量253

点赞数

分类专栏： Image Models

本文链接：https://blog.csdn.net/hahadelaochao/article/details/113453125

版权

Image Models 专栏收录该内容

10 篇文章 0 订阅

订阅专栏

本文介绍了两种注意力机制：空间注意力Spartial_Attention和通道注意力Channel_Attention。空间注意力通过平均池化和最大池化结合卷积层来关注图像的不同区域；通道注意力则利用全局平均池化和全局最大池化后，通过全连接层来捕获不同通道之间的依赖关系，提高特征选择能力。这两种机制可以提升模型的表征学习能力。

摘要由CSDN通过智能技术生成

在这里插入图片描述

class Spartial_Attention(nn.Module):

    def __init__(self, kernel_size):
        super(Spartial_Attention, self).__init__()

        assert kernel_size % 2 == 1, "kernel_size = {}".format(kernel_size)
        padding = (kernel_size - 1) // 2

        self.__layer = nn.Sequential(
            nn.Conv2d(2, 1, kernel_size=kernel_size, padding=padding),
            nn.Sigmoid(),
        )

    def forward(self, x):
        avg_mask = torch.mean(x, dim=1, keepdim=True)
        max_mask, _ = torch.max(x, dim=1, keepdim=True)
        mask = torch.cat([avg_mask, max_mask], dim=1)

        mask = self.__layer(mask)
        return x * mask

class Channel_Attention(nn.Module):

    def __init__(self, channel, r):
        super(Channel_Attention, self).__init__()

        self.__avg_pool = nn.AdaptiveAvgPool2d((1, 1))
        self.__max_pool = nn.AdaptiveMaxPool2d((1, 1))

        self.__fc = nn.Sequential(
            nn.Conv2d(channel, channel//r, 1, bias=False),
            nn.ReLU(True),
            nn.Conv2d(channel//r, channel, 1, bias=False),
        )
        self.__sigmoid = nn.Sigmoid()


    def forward(self, x):
        y1 = self.__avg_pool(x)
        y1 = self.__fc(y1)

        y2 = self.__max_pool(x)
        y2 = self.__fc(y2)

        y = self.__sigmoid(y1+y2)
        return x * y