torch.nn.Multiheadattention介绍

初始化参数:

class
torch.nn.MultiheadAttention(embed_dim, num_heads, dropout=0.0, bias=True, 
add_bias_kv=False, add_zero_attn=False, kdim=None, vdim=None, batch_first=False, 
device=None, dtype=None)

        允许模型共同关注来自不同表示子空间的信息。

多头注意力被定义为:

 如果满足以下所有条件,forward() 将使用 FlashAttention 中描述的优化实现:

         (1)计算自注意力(即查询、键和值是相同的张量。将来会放宽此限制。)

        (2)使用 batch_first==True 对输入进行批处理 (3维)

        (3)autograd 被禁用(使用 torch.inference_mode 或 torch.no_grad)或没有张量参数 requires_grad

        (4)训练被禁用(使用 .eval())

        (5)add_bias_kv 为假

        (6)add_zero_attn 为假

        (7)batch_first 为 True,输入被批处理

        (8)kdim 和 vdim 等于 embed_dim

        (9)如果传递了 NestedTensor,则既不传递 key_padding_mask 也不传递 attn_mask

        (10)autograd被禁用

参数介绍:

        embed_dim – 模型的总维度;

        num_heads——并行注意力头的数量。请注意,embed_dim 将拆分为 num_heads(即每个头都有维度 embed_dim // num_heads);

        dropout——attn_output_weights 上的辍学概率。默认值:0.0(no dropout);

        bias – 如果指定,将偏置添加到输入/输出投影层。默认值:真;

        add_bias_kv – 如果指定,则在 dim=0 处向键和值序列添加偏差。默认值:假;

        add_zero_attn – 如果指定,则将一批新的零添加到 dim=1 处的键和值序列。默认值:假;

        kdim – 键的特征总数。默认值:无(使用 kdim=embed_dim)。

        vdim – 值的特征总数。默认值:无(使用 vdim=embed_dim)。

        batch_first – 如果为 True,则输入和输出张量作为(batch、seq、feature)提供。默认值:False(序列、批次、特征)。

使用案例:

multihead_attn = nn.MultiheadAttention(embed_dim, num_heads)
attn_output, attn_output_weights = multihead_attn(query, key, value)

forward函数

forward(query, key, value, key_padding_mask=None, need_weights=True, 
attn_mask=None, average_attn_weights=True, is_causal=False)

参数解释:

        query (Tensor) – 查询嵌入形状(L,Eq)用于非批处理输入,(L,N,Eq)用于batchfirst= False 或(N,L,Eq)用于batchfirst = True,其中 L是目标序列长度,N 是批处理大小,Eq 是嵌入查询的嵌入维度 embed _ dim。将查询与键-值对进行比较以生成输出。

        key (Tensor)- 非批量输入的形状(S,Ek)的key嵌入,(S,N,Ek)的键嵌入,batchfirst = False 或(N,S,Ek)的键嵌入,batchfirst = True,其中 S 是源序列长度,N 是批量大小,Ek 是key嵌入维度 kdim。

        value (Tensor) – 非批处理输入的形状 (S,Ev) 的value嵌入,(S,N,Ev) 当 batch_first=False 或 (N,S,Ev)当 batch_first=True 时,其中 S 是源序列长度,N是批量大小,Ev 是value嵌入维度 vdim。

        key_padding_mask (Optional[Tensor]) – 如果指定,形状为 (N,S) 的掩码指示要忽略key中的哪些元素以引起注意(即视为“填充”)。对于未批量查询,形状应为 (S)。支持二进制和浮点掩码。对于二进制掩码,True 值表示为了引起注意,相应的key值将被忽略。对于float掩码,它会直接添加到相应的key值。

        need_weights (bool) – 如果指定,除了 attn_outputs 之外还返回 attn_output_weights。默认值:真。

        is_causal (bool) – 如果指定,则应用因果掩码作为注意掩码。与提供 attn_mask 互斥。默认值:假。

        average_attn_weights (bool) – 如果为真,则表示返回的 attn_weights 应该在头部之间进行平均。否则,每个头单独提供 attn_weights。请注意,此标志仅在 need_weights=True 时有效。默认值:True(即头部的平均)

输出:

        attn _ output-当输入未批处理时形状为(L,E),(L,N,E)(L,N,E)当batchfirst = False 或者(N,L,E)当batchfirst= True 时,形状(L,E)(N,L,E)的注意输出,其中 L是目标序列长度,N是批处理大小,E是嵌入维度 embed _ dim。

        attn_output_weights - 仅在 need_weights=True 时返回。如果 average_attn_weights=True,当输入未分批时返回形状为 (L,S),有分批返回(N,L,S)的头部的平均注意力权重,其中 N 是批量大小, L 是目标序列长度,SS 是源序列长度。如果 average_attn_weights=False,当输入未分批或 (N,num_heads,L,S)(N,num_heads,L,S) 时,返回形状为 (num_heads,L,S)(N,num_heads,L,S) 的每个头部的注意力权重。

猜你喜欢

转载自blog.csdn.net/qq_40671063/article/details/130285398