门控注意力机制引入了门控单元,通过门控信号来控制注意力的分配。门控单元可以根据输入的特征动态地调整注意力权重,使得模型能够更加灵活地关注不同的特征或区域。常见的门控单元有 Sigmoid 门控和 Tanh 门控等。

import torch
import torch.nn as nn

class GatedAttention(nn.Module):
    def __init__(self, in_channels):
        super(GatedAttention, self).__init__()
        self.gate = nn.Conv2d(in_channels, 1, kernel_size=1)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        gate = self.gate(x)
        gate = self.sigmoid(gate)
        out = x * gate
        return out

Logo

有“AI”的1024 = 2048,欢迎大家加入2048 AI社区

更多推荐