Pytorch通道注意力和空间注意力代码

请依据自身需要进行更改：代码可以直接运行

import torch
import torch.nn as nn


class ChannelAttention(nn.Module):
    def __init__(self, in_channel, hidden_channel, out_channel):
        super(ChannelAttention, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.max_pool = nn.AdaptiveMaxPool2d(1)
        self.add_pooling = lambda x: self.avg_pool(x)+self.max_pool(x)
        self.conv1 = nn.Conv2d(in_channels=in_channel, out_channels=hidden_channel, kernel_size=(1, 1))   # 依据实际情况可以调整为 线性层或者一维卷积,下同
        self.conv2 = nn.Conv2d(in_channels=hidden_channel, out_channels=out_channel, kernel_size=(1, 1))
        self.sigmoid = nn.Sigmoid()
        self.relu = nn.ReLU()

    def forward(self, x):
        x_pre = self.add_pooling(x)  # (64, 32, 1, 1)
        att = self.sigmoid(self.conv2(self.relu(self.conv1(x_pre))))  # (64, 32, 1, 1)
        x = att*x
        return x


class SpatialAttention(nn.Module):
    def __init__(self):
        super(SpatialAttention, self).__init__()
        self.max_pooling = lambda x: torch.max(x, dim=1, keepdim=True)[0]
        self.ave_pooling = lambda x: torch.mean(x, dim=1, keepdim=True)
        self.conv = nn.Conv2d(2, 1, kernel_size=1)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        apa_att = self.sigmoid(self.conv(torch.cat([self.max_pooling(x), self.ave_pooling(x)], dim=1)))  # (64, 1, 18, 100)
        x = apa_att*x
        return x


if __name__ == "__main__":
    cha_net = ChannelAttention(32, 128, 32)
    spa_net = SpatialAttention()
    data = torch.rand(64, 32, 18, 100) # 自定义一组输入数据（batch, channel, width, length）

    # 两者连用就是CBAM机制，可以根据自身研究需要进行更改，以适应自己的项目
    new_data = cha_net(data)  # (64, 32, 18, 100)
    new_data = spa_net(new_data)  # (64, 32, 18, 100)
    print(new_data.shape)