transfomer中Multi-Head Attention的源码实现

本文主要是介绍transfomer中Multi-Head Attention的源码实现，希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！

简介

Multi-Head Attention是一种注意力机制,是transfomer的核心机制,就是图中黄色框内的部分.

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

源码实现:

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

import torch
from torch import nnclass 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_dimself.num_heads = num_headsself.att_dropout = nn.Dropout(att_dropout)self.out_dropout = nn.Dropout(out_dropout)self.average_attn_weights = average_attn_weightsself.head_dim = embed_dim // num_headsself.scale = self.head_dim ** 0.5assert 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: 参考pytorchReturns:'''assert query.dim() == 3 and key.dim() == 3 and value.dim() == 3assert 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_posif 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 Attentionattention = query @ key.transpose(-2, -1)attention /= self.scale  # 参考: https://blog.csdn.net/zwhdldz/article/details/135462127attention = 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)# shortcutoutput = 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, attentionif __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我在内部加输入映射,具体作用参考下一篇博客。

这篇关于transfomer中Multi-Head Attention的源码实现的文章就介绍到这儿，希望我们推荐的文章对编程师们有所帮助！

---