在Python中如何利用PyTorch实现类似Transformer的自注意力机制？-Python教程-PHP中文网

直接使用pytorch官方nn.multiheadattention即可实现标准multi-head self-attention，它内置缩放点积、mask处理、形状校验和数值稳定机制；需确保输入形状为(seq_len, batch_size, embed_dim)（默认）或启用batch_first=true时用(batch_size, seq_len, embed_dim)，且embed_dim能被num_heads整除，配合正确设备与dtype的attn_mask（如causal mask用torch.triu(..., diagonal=1)填充float('-inf')）可安全支持因果注意力与训练稳定性。

在python中如何利用pytorch实现类似transformer的自注意力机制？

PyTorch里怎么手写一个标准的Multi-Head Self-Attention？

PyTorch不提供开箱即用的独立SelfAttention模块，但nn.MultiheadAttention就是官方实现的自注意力核心——它默认支持mask、batch_first、dropout等关键特性，且底层已做优化。别自己从torch.bmm开始拼QKV，容易漏掉缩放、mask广播或梯度细节。

nn.MultiheadAttention要求输入形状为(seq_len, batch_size, embed_dim)（非batch_first模式），这是最容易出错的地方：传入(batch_size, seq_len, embed_dim)会直接报RuntimeError: expected 3D input
注意embed_dim必须能被num_heads整除，否则初始化就抛ValueError: embed_dim must be divisible by num_heads
如果要用batch_first=True，记得同时设置attn_mask和key_padding_mask的shape匹配：前者是(batch_size, seq_len, seq_len)，后者是(batch_size, seq_len)

attn = nn.MultiheadAttention(embed_dim=512, num_heads=8, batch_first=True)
x = torch.randn(4, 10, 512)  # (batch, seq, dim)
out, _ = attn(x, x, x)  # q, k, v 都是x

为什么直接调`nn.MultiheadAttention`比自己写`torch.einsum`更稳？

自己写QKV计算看似灵活，但实际踩坑密集：比如scale因子忘了除sqrt(d_k)会导致softmax输出趋近one-hot；mask没用torch.where或masked<em>fill</em>处理，会在训练中悄悄引入nan；还有attn_weights @ v维度错位导致matmul失败。

官方nn.MultiheadAttention内部自动做q @ k.transpose(-2, -1) / sqrt(d_k)，且对attn_mask做了float('-inf')填充，行为可预测
它还内置了_mha_shape_check校验逻辑，提前暴露shape问题，而不是等到backward才崩
若需定制（如相对位置编码），建议继承nn.MultiheadAttention重写forward，而非完全重写整个attention计算流

如何让`nn.MultiheadAttention`支持 causal attention（仅看前面token）？

causal mask不是靠参数开关，而是靠传入正确的attn_mask。PyTorch不提供现成的causal_mask函数，得自己生成：

Claw Code Suite (Python-Only Edition)

仅使用 Python 将 Claw Code 虚拟化框架与 OpenClaw 集成，提供 184 种工具和 200 多条命令，用于安全分析和代码质量...

下载

用torch.triu(torch.full((L, L), float('-inf')), diagonal=1)构造上三角mask（注意是float('-inf')，不是0或-1e9，否则fp16下可能失效）
如果用了batch_first=True，mask要扩展成(batch_size, L, L)，不能只传(L, L)，否则广播错误
记得把mask设备设成和输入一致：mask = mask.to(x.device)，否则混合设备会报错

L = x.size(1)
causal_mask = torch.triu(torch.full((L, L), float('-inf')), diagonal=1)
causal_mask = causal_mask.unsqueeze(0)  # (1, L, L)
out, _ = attn(x, x, x, attn_mask=causal_mask)

训练时`nn.MultiheadAttention`突然NaN，大概率是哪几处没设好？

NaN几乎都来自attention权重里的inf或nan参与了后续计算，常见源头：

attn_mask用了0或-1e9代替float('-inf')，在fp16下-1e9可能被截断成正常数，导致softmax分母爆炸
输入x本身含nan（比如前层LayerNorm未设eps，除零后产生NaN），而nn.MultiheadAttention不做输入检查
dropout设得过高（如0.5以上）+ 小batch size，导致某次随机mask后某个head的k全为0，q @ k.T出现全零行，softmax后0/0
使用key_padding_mask时，传入了bool类型mask但没转float，PyTorch 1.12+会静默转换，旧版本可能出错

真正难调试的是mask和输入设备/类型不一致——比如CPU mask + GPU tensor，或者int64 mask传进去，这些不会立刻报错，但会在某些CUDA kernel里触发undefined behavior。每次加attention层，先用.to(dtype=torch.float32).to(device='cuda')统一输入，比事后debug快得多。