neighborhood attention代码实现

邻域注意力（Neighborhood Attention）是一种用于图像分割的注意力机制，可以在局部区域内聚焦于感兴趣的特征。其代码实现如下：

import torch
import torch.nn as nn

class NeighborhoodAttention(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1):
        super(NeighborhoodAttention, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding)
        self.conv2 = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding)
        self.conv3 = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        x1 = self.conv1(x)
        x2 = self.conv2(x)
        x3 = self.conv3(x)
        x4 = torch.cat([x1, x2, x3], dim=1)
        x5 = torch.mean(x4, dim=1, keepdim=True)
        x6 = self.sigmoid(x5)
        x7 = x * x6
        return x7

在这个实现中，我们使用了三个卷积层来提取特征，然后将它们拼接在一起。接着，我们对拼接后的特征进行平均池化，并使用 Sigmoid 函数将其映射到 [0, 1] 的范围内。最后，我们将输入特征和注意力特征相乘，得到最终的输出特征。