优先调用torch.nn.multiheadattention，它已封装q/k/v投影、mask处理等标准流程；需确保embed_dim能被num_heads整除，输入默认(seq_len,batch_size,embed_dim)，设batch_first=true可改为(batch_size,seq_len,embed_dim)。
PyTorch里怎么手写MultiheadAttention模块？
直接用torch.nn.MultiheadAttention就行，不用自己从零实现Q/K/V投影和mask逻辑——它已经封装好了标准流程，且支持batch_first=True、梯度检查、CUDA优化。自己重写容易漏掉attn_mask的广播规则或key_padding_mask的布尔掩码处理，反而引入bug。
实操建议：
- 优先调用
torch.nn.MultiheadAttention(embed_dim, num_heads)，其中embed_dim必须能被num_heads整除 - 输入形状默认是
(seq_len, batch_size, embed_dim)；设batch_first=True可改为(batch_size, seq_len, embed_dim) - 若需自定义初始化，用
reset_parameters()或手动对in_proj_weight赋值，但注意其内部把Q/K/V权重拼成一个大矩阵
为什么forward返回三个张量？
torch.nn.MultiheadAttention.forward()返回(attn_output, attn_output_weights, attn_output_weights)（第三个是冗余的，实际只用前两个），其中attn_output_weights是平均后的注意力权重（shape为(batch_size, num_heads, seq_len, seq_len)），不是原始softmax输出——它做了mean(dim=1)，丢失了头间差异。
如果你要分析单个头的注意力分布：
- 传入
need_weights=True（默认True） - 接住返回值第三项其实是
attn_output_weights的原始未平均版本，但官方API不直接暴露；正确方式是继承并重写forward，在_scaled_dot_product_attention后保留attn张量 - 或者改用
F.scaled_dot_product_attention（PyTorch 2.0+），它返回原始attentions（当is_causal=False且dropout_p=0时）
key_padding_mask和attn_mask的区别在哪？
这是最容易混淆的两个mask参数：key_padding_mask用于屏蔽padding位置（比如句子末尾补零），是(batch_size, seq_len)的布尔张量；attn_mask用于控制任意位置间能否attend（如上三角mask实现因果性），是(seq_len, seq_len)或(batch_size * num_heads, seq_len, seq_len)的float张量（-inf表示屏蔽）。
关键细节：
-
key_padding_mask会被自动广播到每个head，并与attn_mask相加（数值上），所以二者可同时使用 - 若
attn_mask是2D，它会作用于所有batch样本；若想每条样本不同，必须扩展为3D：(batch_size, seq_len, seq_len) - 传
None给任一mask不会报错，但效果不同：不传key_padding_mask就无padding屏蔽；不传attn_mask就无结构化限制
训练时dropout没生效？
默认dropout=0.0，且该dropout只作用于注意力权重（即softmax之后、加权求和之前），不影响Q/K/V投影。如果你发现注意力图完全一致，先确认是否设置了dropout > 0且模型处于train()模式。
还要注意：
-
dropout参数在MultiheadAttention构造时传入，不能在forward中动态改 - PyTorch 2.0+启用
F.scaled_dot_product_attention后，dropout行为更稳定；旧版可能因cudnn路径绕过dropout - 调试时可用
torch.manual_seed()固定，但别在生产中依赖确定性——dropout本就是随机的
真正麻烦的是多卡DDP下attn_mask的设备一致性：如果mask在CPU上而输入在GPU上，会静默失败。务必确保所有tensor同设备，尤其是mask类输入。
Python免费学习笔记（深入）：立即使用
在学习笔记中，你将探索 Python 的核心概念和高级技巧！
