PyTorch nn.MultiHead() 参数理解

最新推荐文章于 2025-05-19 13:37:30 发布

springtostring

最新推荐文章于 2025-05-19 13:37:30 发布

阅读量1.1w

点赞数 10

CC 4.0 BY-SA版权

分类专栏：计算机编程

本文链接：https://blog.youkuaiyun.com/springtostring/article/details/113958933

计算机编程专栏收录该内容

17 篇文章

订阅专栏

本文详细介绍了PyTorch中nn.MultiheadAttention模块的使用方法和注意事项，包括模型初始化、参数解释、输入输出格式，并通过实例展示了如何避免报错。讨论了embed_dim与num_heads的关系，强调了它们之间的约束条件。此外，还探讨了attn_mask参数的作用，用于在self-attention中屏蔽不可见的位置。

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

之前一直是自己实现MultiHead Self-Attention程序，代码段又臭又长。后来发现Pytorch 早已经有API nn.MultiHead()函数，但是使用时我却遇到了很大的麻烦。

首先放上官网说明：
$MultiHead(Q,K,V)=Concat(head_1,…,head_h)W_O\quad where\ head_i=Attention(QW_i^Q,KW_i^K,VW_i^V)$

# 模型初始化
torch.nn.MultiheadAttention(embed_dim, num_heads, dropout=0.0, bias=True, add_bias_kv=False, add_zero_attn=False, kdim=None, vdim=None)
'''
embed_dim – 嵌入向量总长度.

num_heads – 并行的head数目，即同时做多少次不同语义的attention.

dropout – dropout的概率.

bias – 是否添加偏置.默认: True.
'''


# 模型运算
forward(query, key, value, key_padding_mask=None, need_weights=True, attn_mask=None)
'''
Inputs:

query: (L, N, E) where L is the target sequence length, N is the batch size, E is the embedding dimension.

key: (S, N, E) , where S is the source sequence length, N is the batch size, E is the embedding dimension.

value: (S, N, E) where S is the source sequence length, N is the batch size, E is the embedding dimension.

key_padding_mask: (N, S)(N,S) , ByteTensor, where N is the batch size, S is the source sequence length.

attn_mask: 2D mask (L, S)(L,S) where L is the target sequence length, S is the source sequence length. 3D mask (N*num_heads, L, S)(N∗num_heads,L,S) where N is the batch size, L is the target sequence length, S is the source sequence length.

Outputs:

attn_output: (L, N, E)(L,N,E) where L is the target sequence length, N is the batch size, E is the embedding dimension.

attn_output_weights: (N, L, S)(N,L,S) where N is the batch size, L is the target sequence length, S is the source sequence length.
'''

值得注意的一点是，query，key，value的输入形状一定是 [sequence_size, batch_size, emb_size]

官网例子：

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

embed_dim, num_heads参数

但是我执行程序却报错了：

A = torch.arange(1,25).view(4,3,2)
A = A.float()

self_attn = torch.nn.MultiheadAttention(embed_dim=2, num_heads=4, dropout=0.0)
res,weight = self_attn(A,A,A)

报错信息：

    self_attn = torch.nn.MultiheadAttention(embed_dim=2, num_heads=4, dropout=0.0)
  File "E:\Anaconda3\lib\site-packages\torch\nn\modules\activation.py", line 740, in __init__
    assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads"
AssertionError: embed_dim must be divisible by num_heads

切到源码看不懂，而且我用pycharm 一直ctrl+鼠标进入不了最底层的代码，只有前几层的代码。(求相关领域大佬教教我)

经过自己的尝试，nn.MultiheadAttention(embed_dim, num_heads)中的要满足两点约束：

embed_dim == input_dim ，即query，key，value的embedding_size必须等于embed_dim
embed_dim%num_heads==0

上面的约束，也就说明了在使用nn.MultiheadAttention(embed_dim, num_heads)时， num_heads不是我们想设多少就设定多少。

我的看法：

nn.MultiheadAttention(embed_dim, num_heads) 中的embed_dim 是输入的embeddingsize，即query输入形状(L, N, E)的E数值，nn.MultiheadAttention 想要实现的是无论head 数目设置成多少，输出的向量大小都是不变的。写完这句话，发现自己表达能力真是弱，自己都看不懂。可以结合下面公式理解。
$MultiHead(Q,K,V)=Concat(head_1,…,head_h)W_O\quad where\ head_i=Attention(QW_i^Q,KW_i^K,VW_i^V)\\ Q\in R^{(L, N, E)},\ K\in R^{(S, N, E)},\ V\in R^{(S, N, E)},\ W_i\in R^{(E,E/h)},W_O\in R^{(E,E)}$

attn_mask参数

self-attention公式: $SA⁡(Q,K,V)=softmax(QKTdk)V\operatorname{SA}(Q, K, V)=softmax\left(\frac{Q K^{T}}{\sqrt{d_{k}}}\right) V$

$Q K^{T}$ 生成权重分布，但是在应用中有些位置的权重是不可见的，比如在时间序列中，第t天时，我们并不知道t+1天之后的信息。这时就需要传入attn_mask参数，屏蔽这些不合理的权重。attn_mask要求是booltensor，某个位置true表示掩盖该位置。

import torch
import torch.nn as nn

A = torch.Tensor(5,2,4)
nn.init.xavier_normal_(A)
print(A)

# tensor([[[ 0.3688,  0.0391,  0.2048, -0.0906],
#          [-0.0654,  0.1193, -0.1792,  0.0470]],
#
#         [[ 0.0812, -0.4180, -0.1353, -0.2670],
#          [ 0.0433,  0.1442,  0.1733,  0.0535]],
#
#         [[ 0.2352, -0.3314, -0.0238,  0.4116],
#          [ 0.1062,  0.5122,  0.1572, -0.2991]],
#
#         [[ 0.3381,  0.4004, -0.1936, -0.1553],
#          [-0.0168,  0.5914,  0.7389, -0.1740]],
#
#         [[ 0.0446, -0.1739, -0.2020,  0.2580],
#          [-0.0109,  0.0854,  0.2634, -0.4735]]])

M = nn.MultiheadAttention(embed_dim=4, num_heads=2)
attention_mask = ~torch.tril(torch.ones([A.shape[0],A.shape[0]])).bool()
print(attention_mask)

# tensor([[False,  True,  True,  True,  True],
#         [False, False,  True,  True,  True],
#         [False, False, False,  True,  True],
#         [False, False, False, False,  True],
#         [False, False, False, False, False]])

attn_output, attn_output_weights=M(A,A,A, attn_mask=attention_mask)
print(attention_mask)

# tensor([[[1.0000, 0.0000, 0.0000, 0.0000, 0.0000],
#          [0.5067, 0.4933, 0.0000, 0.0000, 0.0000],
#          [0.3350, 0.3276, 0.3374, 0.0000, 0.0000],
#          [0.2523, 0.2511, 0.2549, 0.2417, 0.0000],
#          [0.2004, 0.1962, 0.2039, 0.1981, 0.2013]],
# 
#         [[1.0000, 0.0000, 0.0000, 0.0000, 0.0000],
#          [0.5025, 0.4975, 0.0000, 0.0000, 0.0000],
#          [0.3325, 0.3312, 0.3363, 0.0000, 0.0000],
#          [0.2535, 0.2429, 0.2633, 0.2404, 0.0000],
#          [0.2002, 0.1986, 0.2008, 0.1976, 0.2028]]], grad_fn=<DivBackward0>)