门控注意力机制（Gated Attention）

wuhansong0592

于 2025-02-24 10:56:27 发布

阅读量1k

点赞数 3

文章标签：深度学习 python

本文链接：https://blog.youkuaiyun.com/wuhansong0592/article/details/145823159

版权

门控注意力机制引入了门控单元，通过门控信号来控制注意力的分配。门控单元可以根据输入的特征动态地调整注意力权重，使得模型能够更加灵活地关注不同的特征或区域。常见的门控单元有 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