手写SelfAttention层需四步：线性投影得Q/K/V；缩放点积Q@Kᵀ/√dₖ；mask填-inf后softmax；加权求和再经Wₒ输出，严格控维、控dtype、控mask位置。
PyTorch里怎么手写SelfAttention层（不依赖nn.MultiheadAttention）
直接实现最简自注意力，能帮你真正理解Q、K、V的计算逻辑和缩放因子的作用。别急着套用封装好的模块——很多 bug 出在维度没对齐或mask应用位置错。
关键点：输入x形状必须是(batch_size, seq_len, embed_dim)；W_q、W_k、W_v都是(embed_dim, head_dim)；算Q @ K.T前要除以sqrt(head_dim)，否则 softmax 容易饱和。
- 用
torch.bmm做 batch 内矩阵乘（不能直接@，维度会错） - 如果要做 causal attention，mask 要在 softmax 前加，用
torch.finfo(torch.float32).min填无效位置 - 输出要再过一个线性层
W_o（(head_dim, embed_dim)），不然维度接不上后续层
class SelfAttention(torch.nn.Module):
def __init__(self, embed_dim, head_dim):
super().__init__()
self.W_q = torch.nn.Linear(embed_dim, head_dim, bias=False)
self.W_k = torch.nn.Linear(embed_dim, head_dim, bias=False)
self.W_v = torch.nn.Linear(embed_dim, head_dim, bias=False)
self.W_o = torch.nn.Linear(head_dim, embed_dim, bias=False)
self.scale = head_dim ** 0.5
<pre class='brush:python;toolbar:false;'>def forward(self, x, mask=None):
Q = self.W_q(x) # (b, s, d)
K = self.W_k(x) # (b, s, d)
V = self.W_v(x) # (b, s, d)
attn = torch.bmm(Q, K.transpose(1, 2)) / self.scale # (b, s, s)
if mask is not None:
attn = attn.masked_fill(mask == 0, float('-inf'))
attn = torch.softmax(attn, dim=-1)
out = torch.bmm(attn, V) # (b, s, d)
return self.W_o(out)为什么nn.MultiheadAttention默认要求batch_first=False
因为 PyTorch 的历史设计沿用了 NLP 传统：把seq_len放在第 0 维（即(seq_len, batch_size, embed_dim)）。如果你强行传入batch_first=True，内部会自动 transpose，但容易和你自己写的 layer 维度混掉——尤其接在nn.TransformerEncoderLayer后面时，输出默认还是batch_first=False格式。
- 检查输入 shape：若你习惯
(B, S, E)，就显式设batch_first=True，且确保所有前后层统一 - 注意
key_padding_mask形状是(B, S)，而attn_mask（用于 causal 或自定义遮罩）是(S, S)或(B, S, S) - 该模块不自动做 LayerNorm 或残差连接，得自己包一层
TransformerEncoderLayer才完整
常见报错：RuntimeError: mat1 and mat2 shapes cannot be multiplied
90% 出现在Q @ K.T这步。不是模型写错了，而是输入x的embed_dim和线性层的in_features不一致，或者batch_first开关和实际数据 shape 冲突。
立即学习“Python免费学习笔记（深入）”；
- 打印
x.shape和self.W_q.weight.shape，确认第二维匹配 - 如果用了
nn.Embedding，注意它输出是(B, S, E)，但 embedding 层的num_embeddings和embedding_dim别和 attention 的embed_dim搞混 - GPU 上调试时，
torch.cuda.synchronize()后加断点，避免异步报错定位偏移
要不要自己实现，还是直接用nn.MultiheadAttention
除非你在调参、做消融实验，或需要定制梯度（比如让attn_weights参与 loss），否则别手写。PyTorch 版本 ≥1.10 后，nn.MultiheadAttention已支持 FlashAttention 风格优化（通过enable_math=False等 flag），速度比手写快，且自动处理了多头拆分、concat、dropout 等细节。
- 想兼容旧版 PyTorch（
- 训练时发现 attention 权重异常（全 0 或全 1），优先检查
scale因子和mask是否生效，而不是重写整个模块 - 导出 ONNX 时，手写 attention 更容易适配，但
nn.MultiheadAttention从 1.12 起也基本稳定了
维度对齐永远比公式推导更耗时间，尤其是mask的广播行为和bmm的 batch 维位置——这些地方错一点，结果就全崩。
