PyTorch 中的 `nn.MultiheadAttention` 模块实现了多头注意力机制（Multi-Head Attention），这是 Transformer 架构的核心组成部分。下面是一个详细的使用说明： ### 基本概念 在 PyTorch 的 Multihead Attention 实现中，输入主要包括三个部分：查询(Query)、键(Key) 和值(Value)。此外还需要指定嵌入维度(embed_dim)及头部数(num_heads)等相关参数。 --- ### 使用示例 ```python import torch import torch.nn as nn # 初始化 multihead_attention 层 embed_size = 128 # 输入特征向量的维数 num_heads = 8 # 头的数量 mha = nn.MultiheadAttention(embed_dim=embed_size, num_heads=num_heads) # 创建随机张量模拟 Query, Key 和 Value 数据 batch_size = 32 # 批次大小 seq_length = 50 # 序列长度 query = torch.randn(seq_length, batch_size, embed_size) key = torch.randn(seq_length, batch_size, embed_size) value = torch.randn(seq_length, batch_size, embed_size) # 运行 multi-head attention 并获取结果 output, attn_weights = mha(query=query, key=key, value=value) print("Output shape:", output.shape) # 输出形状 [sequence length, batch size, embedding dim] print("Attention weights shape:", attn_weights.shape) # 注意力权重 [batch size * number of heads, sequence length, sequence length] ``` 在这个例子中，我们初始化了一个拥有八个头的多头注意层，并生成了一些随机数据作为查询、键和值来进行测试。 --- ### 参数解释 1. **embed_dim**: 表征每个token的特征向量的总尺寸。 2. **num_heads**: 将嵌入空间划分为多少个独立子空间去单独计算attention score. 3. **dropout**: 可选，默认无。用于防止过拟合，在forward pass过程中按照一定比率drop一些attentionscores。 4. **bias**: 默认True。是否添加偏差项给Q,K,V转换矩阵W_q,W_k,W_v上。 5. **add_zero_attn**: 如果设为true，则会在keys/values前面加上一组全零的数据点供第一个时间步长参考用。（较少用） 6. **device & dtype**: 定义使用的设备(CPU/GPU)以及数据类型(float,double etc.) 7. 其他还有类似out_proj_weight/bias这样的细节点可以调整但一般很少直接设置除非特殊需求下才修改默认配置. --- ### 返回值解析 - **output**: 经过多头注意力操作后的输出张量，其shape与input相同 `[L,B,E]`. - **attn_output_weights**: 形状为 `(B*num_head,L,S)` 的tensor表示的是每一对query-key之间的相似度得分经过softmax之后形成的probability distribution.