pytorch实现mutil_head attention

最新推荐文章于 2025-03-21 16:10:46 发布

经十东路31699号

最新推荐文章于 2025-03-21 16:10:46 发布

阅读量6.1k

点赞数 10

分类专栏：模型与方法文章标签：自然语言处理

本文链接：https://blog.csdn.net/qq_37262929/article/details/110056809

版权

该博客介绍了如何使用PyTorch实现点积注意力和多头注意力机制。首先，展示了点积注意力的实现，包括softmax和dropout操作。接着，详细解释了多头注意力的实现，涉及线性映射、头的分割和合并，以及最终的LayerNorm和dropout。代码示例中，展示了如何应用这些机制处理输入张量。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

点积注意力的实现方法

import torch
import torch.nn as nn
import numpy as np


class dot_attention(nn.Module):
    """ 点积注意力机制"""

    def __init__(self, attention_dropout=0.0):
        super(dot_attention, self).__init__()
        self.dropout = nn.Dropout(attention_dropout)
        self.softmax = nn.Softmax(dim=2)

    def forward(self, q, k, v, scale=None, attn_mask=None):
        """
        前向传播
        :param q:
        :param k:
        :param v:
        :param scale:
        :param attn_mask:
        :return: 上下文张量和attention张量。
        """
        attention = torch.bmm(q, k.transpose(1, 2))
        if scale:
            attention = attention * scale        # 是否设置缩放
        if attn_mask:</