transfomer中Multi-Head Attention的源码实现

这篇具有很好参考价值的文章主要介绍了transfomer中Multi-Head Attention的源码实现。希望对大家有所帮助。如果存在错误或未考虑完全的地方,请大家不吝赐教,您也可以点击"举报违法"按钮提交疑问。

本专栏主要是深度学习/自动驾驶相关的源码实现,获取全套代码请参考

简介

Multi-Head Attention是一种注意力机制,是transfomer的核心机制,就是图中黄色框内的部分.
transfomer中Multi-Head Attention的源码实现,手撕源码系列,BEV感知系列,深度学习,人工智能,pytorch,transformer,自然语言处理,神经网络,目标检测

Multi-Head Attention的原理是通过将模型分为多个头,形成多个子空间,让模型关注不同方面的信息。每个头独立进行注意力运算,得到一个注意力权重矩阵。输出的结果再通过线性变换和拼接操作组合在一起。这样可以提高模型的表示能力和泛化性能。
在Multi-Head Attention中,每个头的权重矩阵是随机初始化生成的,并在训练过程中通过梯度下降等优化算法进行更新。通过这种方式,模型可以学习到如何将输入序列的不同部分关联起来,从而捕获更多的上下文信息。
总之,Multi-Head Attention通过将模型分为多个头,形成多个子空间,让模型关注不同方面的信息,提高了模型的表示能力和泛化性能。它的源码实现基于Scaled Dot-Product Attention,通过并行运算和组合输出来实现多头注意力机制。

源码实现:

具体源码及其注释如下,配好环境可直接运行:

import torch
from torch import nn


class MultiheadAttention(nn.Module):
    def __init__(self,
                 embed_dim,
                 num_heads,
                 att_dropout=0.1,
                 out_dropout=0.1,
                 average_attn_weights=True,
                 use_separate_proj_weight = False,
                 device=None,
                 dtype=None):
        super(MultiheadAttention, self).__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.att_dropout = nn.Dropout(att_dropout)
        self.out_dropout = nn.Dropout(out_dropout)
        self.average_attn_weights = average_attn_weights
        self.head_dim = embed_dim // num_heads
        self.scale = self.head_dim ** 0.5
        assert self.embed_dim == self.num_heads * self.head_dim, \
            'embed_dim <{}> must be divisible by num_heads <{}>'.format(self.embed_dim, self.num_heads)
        self.fuse_heads = nn.Linear(self.embed_dim, self.embed_dim)
        factory_kwargs = {'device': device, 'dtype': dtype}
        self.use_separate_proj_weight = use_separate_proj_weight # 是否对输入进行线性映射
        if not use_separate_proj_weight:
            self.in_proj_weight = nn.Parameter(torch.empty((3 * embed_dim, embed_dim), **factory_kwargs))
            self.in_proj_bias = nn.Parameter(torch.empty(3 * embed_dim, **factory_kwargs))
        self._reset_parameters()

    def _reset_parameters(self):
        nn.init.xavier_uniform_(self.in_proj_weight)
        nn.init.constant_(self.in_proj_bias, 0.)

    def forward(self,
                query: torch.Tensor,
                key: torch.Tensor,
                value: torch.Tensor,
                identity=None,
                query_pos=None,
                key_pos=None,
                use_separate_proj_weight: bool = False):
        '''
        Args:
            query:
            key:
            value:
            identity:
            query_pos:
            key_pos:
            use_separate_proj_weight: 参考pytorch

        Returns:

        '''
        assert query.dim() == 3 and key.dim() == 3 and value.dim() == 3
        assert key.shape == value.shape, f"key shape {key.shape} does not match value shape {value.shape}"
        tgt_len, bsz, embed_dim = query.shape  # [查询数量 batch数量 特征维度]
        src_len, _, _ = key.shape  # [被查询数量,_,_]
        # 默认和query进行shortcut(要在位置编码前,因为output为输出特征,特征和原特征shortcut,下一层再重新加位置编码,否则不就重了)
        if identity is None:
            identity = query.clone()
        # 位置编码
        if query_pos is not None:
            query = query + query_pos
        if key_pos is not None:
            key = key + key_pos

        # 是否需要对输入进行映射,mmcv中 q=k=v,那么就需要此处进行映射
        if not self.use_separate_proj_weight:
            assert self.in_proj_weight is not None, "use_separate_proj_weight is False but in_proj_weight is None"
            query, key, value = nn.functional._in_projection_packed(query, key, value, self.in_proj_weight, self.in_proj_bias)
        # 特征划分为self.num_heads 份 [tgt,b,embed_dim] -> [b,n_h, tgt, d_h]
        # [n,b,n_h*d_h] -> [b,n_h,n,d_h] 主要是target和source之前的特征匹配和提取, batch和n_h维度不处理
        query = query.contiguous().view(tgt_len, bsz, self.num_heads, self.head_dim).permute(1, 2, 0, 3)
        key = key.contiguous().view(src_len, bsz, self.num_heads, self.head_dim).permute(1, 2, 0, 3)
        value = value.contiguous().view(src_len, bsz, self.num_heads, self.head_dim).permute(1, 2, 0, 3)
        # [b,n_h,tgt_len,src_len]
        # Scaled Dot-Product Attention
        attention = query @ key.transpose(-2, -1)
        attention /= self.scale  # 参考: https://blog.csdn.net/zwhdldz/article/details/135462127
        attention = torch.softmax(attention, dim=-1)  # 行概率矩阵
        attention = self.att_dropout(input=attention)  # 正则化方法 DropKey,用于缓解 Vision Transformer 中的过拟合问题
        # [b,n_h,tgt_len,d_h] = [b,n_h,tgt_len,src_len] * [b,n_h,src_len,d_h]
        output = attention @ value
        # [b,n_h,tgt_len,d_h] -> [b,tgt_len,embed_dim]
        output = output.permute(0, 2, 1, 3).contiguous().view(tgt_len, bsz, embed_dim)
        # 头之间通过全连接融合一下
        output = self.fuse_heads(output)
        output = self.out_dropout(output)
        # shortcut
        output = output + identity
        # 多头head求平均
        if self.average_attn_weights:
            attention = attention.sum(dim=1) / self.num_heads
        # [tgt_len,b,embed_dim],[b,tgt_len,src_len]
        return output, attention


if __name__ == '__main__':
    query = torch.rand(size=(10, 2, 64))
    key = torch.rand(size=(5, 2, 64))
    value = torch.rand(size=(5, 2, 64))
    query_pos = torch.rand(size=(10, 2, 64))
    key_pos = torch.rand(size=(5, 2, 64))

    att = MultiheadAttention(64, 4)
    # 返回特征采样结果和attention矩阵
    output = att(query=query, key=key, value=value,query_pos=query_pos,key_pos=key_pos)
    pass


具体流程说明:

  1. 将input映射为qkv,如果是cross_attention,q与kv的行数可以不同,但列数(编码维度/通道数)必须相同
  2. q和v附加位置编码
  3. Scaled Dot-Product :通过计算Query和Key之间的点积除以scale得到注意力权重,经过dropout再与Value矩阵相乘得到输出。*scale和dropout的说明参考我的上一篇博客
  4. 输出的结果再通过线性变换融合多头信息。

在实现中,参考pytorch我在内部加输入映射,具体作用参考下一篇博客。

如需获取全套代码请参考文章来源地址https://www.toymoban.com/news/detail-795765.html

到了这里,关于transfomer中Multi-Head Attention的源码实现的文章就介绍完了。如果您还想了解更多内容,请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章,希望大家以后多多支持TOY模板网!

本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如若转载,请注明出处: 如若内容造成侵权/违法违规/事实不符,请点击违法举报进行投诉反馈,一经查实,立即删除!

领支付宝红包 赞助服务器费用

相关文章

  • 学习Transformer前言(Self Attention Multi head self attention)

    一直在做项目,也比较懒没有挤出时间去学习新的东西,感觉停滞很久了,好长一段时间都没有新的知识输入,早就需要就去学习transformer了,因此先来学习注意力机制,本文为个人的一个笔记总结。主要是基于李宏毅老师的一个课程视频笔记,论文原文,加上B站UP主的霹雳

    2024年02月02日
    浏览(54)
  • multi-head_seft-attention(多头自注意力)

    相比于single-head,multi-head就是将 q i q^i q i 分成了 h h h 份 将 q i q^i q i 分成了 h h h 份 计算过程 对于每个Head,我们可以提取出他的 b 11 b_{11} b 11 ​ 到 b m 1 b_{m1} b m 1 ​ ,以 H e a d 1 Head_1 He a d 1 ​ 举例 将输入序列进行embedding后,变为向量 a 1 a_1 a 1 ​ , a 2 a_2 a 2 ​ , a 3 a_3 a 3 ​

    2024年02月13日
    浏览(46)
  • transfomer中Decoder和Encoder的base_layer的源码实现

    本专栏主要是深度学习/自动驾驶相关的源码实现,获取全套代码请参考 Encoder和Decoder共同组成transfomer,分别对应图中左右浅绿色框内的部分. Encoder: 目的:将输入的特征图转换为一系列自注意力的输出。 工作原理:首先,通过卷积神经网络(CNN)提取输入图像的特征。然后,

    2024年01月18日
    浏览(55)
  • Multi-Task Learning based Video Anomaly Detection with Attention 论文阅读

    文章信息: 原文链接:https://ieeexplore.ieee.org/document/10208994/ 源代码:无 发表于:CVPR 2023 基于多任务学习的视频异常检测方法将多个代理任务结合在不同的分支中,以便在不同情境中检测视频异常。然而,大多数现有方法存在以下一些缺点: I) 它们的代理任务组合方式不是以

    2024年01月20日
    浏览(53)
  • Transformer的Q、K、V和Mutil-Head Self-Attention(超详细解读)

    目录 一.什么是Q、K、V 二.Mutil-Head Self-Attention Transformer大行其道,在众多领域取得了不可忽视的成就。如今大火的语言大模型LLM也都是基于Transformer,但是Transformer中的Q、K、V和多头注意力到底是什么呢?这里简单做个学习记录,进行再一次认识和掌握。 Transformer中的Q、K和

    2024年02月06日
    浏览(32)
  • 【论文笔记】Triplet attention and dual-pool contrastive learning for clinic-driven multi-label medical...

    多标签分类Multi-label classification (MLC)可在单张图像上附加多个标签,在医学图像上取得了可喜的成果。但现有的多标签分类方法在实际应用中仍面临着严峻的临床现实挑战,例如: 错误分类带来的医疗风险, 不同疾病之间的样本不平衡问题 无法对未预先定义的疾病(未见疾

    2024年02月03日
    浏览(47)
  • 口语理解任务源码详解系列(二)利用seq2seq-attention模型实现

    写在前面 在前文介绍了项目的数据集构建:传送门 本文利用seq2seq-attention实现:实现细节请参考论文: 《Attention-Based Recurrent Neural Network Models for Joint Intent Detection and Slot Filling》 1.意图识别是利用encoder中的最后一个时间步中的双向隐层 + encoder的attention,最后接一个全连接层

    2024年02月12日
    浏览(52)
  • 【手撕Spring源码】AOP

    AOP 底层实现方式之一是代理,由代理结合通知和目标,提供增强功能 除此以外,aspectj 提供了两种另外的 AOP 底层实现: 第一种是通过 ajc 编译器在 编译 class 类文件时,就把通知的增强功能,织入到目标类的字节码中 第二种是通过 agent 在 加载 目标类时,修改目标类的字节

    2024年02月07日
    浏览(45)
  • 思维导图手撕MyBatis源码

    既然要读MyBatis的源码,那么我们就要先弄清楚MyBatis的入口在哪。这里我们直接写一个标准的MyBatis使用程序,以此来寻找入口分析源码: 这个地方你替换成为这个也是可以的; 或者 这个方法更加通用,在web环境和java环境下均可以使用 这个地方使用了建造者模式。build方法返

    2023年04月15日
    浏览(40)
  • 【手撕Spring源码】深度理解SpringMVC【上】

    既然我们讨论SpringMVC那么就必然绕不开一个东西叫做DispatcherServlet。 DispatcherServlet是SpringMVC的核心Servlet,也叫做前端控制器。它的主要作用是调度请求并将请求分发给相应的处理器。 我们要注意: DispatcherServlet由Servlet容器创建,并且它的生命周期也是Servlet那套体系由Servle

    2024年02月07日
    浏览(41)

觉得文章有用就打赏一下文章作者

支付宝扫一扫打赏

博客赞助

微信扫一扫打赏

请作者喝杯咖啡吧~博客赞助

支付宝扫一扫领取红包,优惠每天领

二维码1

领取红包

二维码2

领红包