ELA高效局部注意力机制（paddle代码实现）

小白一枚，记录学习成果，请指教~

高效局部注意力（ELA）模块，能够准确捕获感兴趣区域的位置，保持输入特征图通道的维度，并保留其轻量级的特性。与CA类似，ELA在空间维度上使用strip pooling获得水平和垂直方向的特征向量。独立处理每个方向上的上述特征向量以获得注意力预测，然后使用乘积运算将它们组合起来，确保感兴趣区域的准确位置信息。具体来说，在第二步中，应用一维卷积分别与两个特征向量进行局部交互，可以选择调整核大小以指示局部交互的范围。得到的特征向量经过分组归一化（GN）和非线性激活函数处理，以产生两个方向上的位置注意力预测。最终的位置注意力是通过将两个方向上的位置注意力相乘得到的。与二维卷积相比，一维卷积更适合处理序列信号，并且更轻量、更快。GN与BN相比表现出相当的性能和更强的泛化能力。

 论文地址：https://arxiv.org/pdf/2403.01123.pdf

# ELA (coordinate attention)

import paddle
import paddle.nn as nn


class ELA(nn.Layer):
    def __init__(self, in_ch, kernel_size=7):
        super(ELA, self).__init__()
        self.pool_h = nn.AdaptiveAvgPool2D((None, 1))
        self.pool_w = nn.AdaptiveAvgPool2D((1, None))

        pad = kernel_size // 2
        self.conv1 = nn.Conv1D(in_ch, in_ch, kernel_size=kernel_size, padding=pad, groups=in_ch, bias_attr=False)
        self.gn = nn.GroupNorm(16, in_ch)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        identity = x
        b, c, h, w = x.shape

        x_h = self.pool_h(x)
        x_w = self.pool_w(x)

        x_h = x_h.reshape([b, c, h])
        x_w = x_w.reshape([b, c, w])

        x_h = self.conv1(x_h)
        x_h = self.gn(x_h)
        x_h = self.sigmoid(x_h).reshape([b, c, h, 1])

        x_w = self.conv1(x_w)
        x_w = self.gn(x_w)
        x_w = self.sigmoid(x_w).reshape([b, c, 1, w])

        out = identity * x_w * x_h

        return out


if __name__ == '__main__':
    ela = ELA(512)  # in_channel
    x = paddle.randn([64, 512, 14, 14])  # (batchsize, channel, H, W)
    y = ela(x)
    print(y.shape)