ECA 注意力机制 代码

您好！下面是一个简单的示例代码，演示了如何实现 ECA 注意力机制：

import torch
import torch.nn as nn

class ECAAttention(nn.Module):
    def __init__(self, in_channels, k_size=3):
        super(ECAAttention, self).__init__()

        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.conv = nn.Conv1d(1, 1, kernel_size=k_size, padding=(k_size - 1) // 2, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        # x shape: (batch_size, in_channels, height, width)
        b, c, h, w = x.size()

        y = self.avg_pool(x)  # y shape: (batch_size, in_channels, 1, 1)
        y = y.view(b, c)  # y shape: (batch_size, in_channels)

        y = y.unsqueeze(1)  # y shape: (batch_size, 1, in_channels)
        y = self.conv(y)  # y shape: (batch_size, 1, in_channels)

        y = y.view(b, 1, c, 1)  # y shape: (batch_size, 1, in_channels, 1)
        y = self.sigmoid(y)  # y shape: (batch_size, 1, in_channels, 1)

        return x * y.expand_as(x)

# 使用示例
in_channels = 64
k_size = 3

input = torch.randn(16, in_channels, 32, 32)
eca = ECAAttention(in_channels, k_size)

output = eca(input)
print(output.shape)

在上述代码中，`ECAAttention` 是一个自定义的注意力模块，它接收输入张量 `x` 并返回经过 ECA 注意力机制处理后的输出张量 `output`。注意力机制主要包括全局平均池化、一维卷积和 Sigmoid 激活函数。通过将 `ECAAttention` 实例化并将输入张量传递给它的 `forward` 方法，我们可以得到注意力机制的输出结果。

希望这个示例代码能对您有所帮助！如果您有任何其他问题，请随时提问。