**PyTorch 极简实现 Self‑Attention 可运行代码（带注释版）-北京尧图网络

下面给出缩放点积自注意力Scaled Dot‑Product Self‑Attention极简 PyTorch 实现全部代码可直接运行注释详细不含多头先实现基础版之后附带多头注意力。import torch import torch.nn as nn import torch.nn.functional as F class ScaledDotProductSelfAttention(nn.Module): def __init__(self, embed_dim): super().__init__() self.d_k embed_dim # Q,K向量维度 # 可训练权重矩阵 Wq, Wk, Wv self.W_q nn.Linear(embed_dim, embed_dim) self.W_k nn.Linear(embed_dim, embed_dim) self.W_v nn.Linear(embed_dim, embed_dim) def forward(self, x, maskNone): :param x: input, shape [batch_size, seq_len, embed_dim] :param mask: 掩码矩阵 [batch_size, seq_len, seq_len] GPT解码器用上三角maskBERT默认None :return out: [batch_size, seq_len, embed_dim] attn_weight: 注意力权重 [batch, seq_len, seq_len] batch, seq_len, _ x.shape # Step1生成 Q, K, V Q self.W_q(x) # (B, L, D) K self.W_k(x) V self.W_v(x) # Step2: Q K^T / sqrt(d_k) attn_score torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.d_k, dtypetorch.float32)) # 如果传入mask将mask为0的位置设置为‑infsoftmax之后权重趋近0 if mask is not None: attn_score attn_score.masked_fill(mask 0, -1e9) # Step3 softmax得到注意力权重 attn_weight F.softmax(attn_score, dim-1) # Step4 加权value得到输出 out torch.matmul(attn_weight, V) return out, attn_weight if __name__ __main__: # 超参设置 batch_size 2 seq_len 4 # 句子4个token embed_dim 64 # 词嵌入维度 # 随机模拟输入向量等价4个单词的embedding x torch.randn(batch_size, seq_len, embed_dim) # 实例化自注意力模块 self_attn ScaledDotProductSelfAttention(embed_dim) output, attn_weights self_attn(x) print(输出shape, output.shape) # torch.Size([2,4,64]) print(注意力权重shape, attn_weights.shape) # torch.Size([2,4,4])2. 拓展多头注意力 Multi‑Head‑AttentionPyTorch 实现class MultiHeadSelfAttention(nn.Module): def __init__(self, embed_dim, num_heads): super().__init__() assert embed_dim % num_heads 0, 嵌入维度必须可以被头数整除 self.d_model embed_dim self.num_heads num_heads self.d_k embed_dim // num_heads self.Wq nn.Linear(embed_dim, embed_dim) self.Wk nn.Linear(embed_dim, embed_dim) self.Wv nn.Linear(embed_dim, embed_dim) self.W_o nn.Linear(embed_dim, embed_dim) def forward(self, x, maskNone): B, L, _ x.shape Q self.Wq(x) K self.Wk(x) V self.Wv(x) # 将维度拆分多头: (B, heads, seq_len, d_k) Q Q.view(B, L, self.num_heads, self.d_k).permute(0,2,1,3) K K.view(B, L, self.num_heads, self.d_k).permute(0,2,1,3) V V.view(B, L, self.num_heads, self.d_k).permute(0,2,1,3) # 缩放点积注意力 attn_score torch.matmul(Q, K.transpose(-2,-1)) / torch.sqrt(torch.tensor(self.d_k, dtypetorch.float32)) if mask is not None: attn_score attn_score.masked_fill(mask 0, -1e9) attn_weight F.softmax(attn_score, dim-1) out torch.matmul(attn_weight, V) # 拼接多头结果 out out.permute(0,2,1,3).contiguous().view(B, L, self.d_model) out self.W_o(out) return out, attn_weight if __name__ __main__: batch 2 seq_len 4 embed_dim 64 heads 8 x torch.randn(batch, seq_len, embed_dim) mh_attn MultiHeadSelfAttention(embed_dim, heads) out, attn mh_attn(x) print(多头注意力输出shape, out.shape)3. 关键点总结Self‑Attention 条件Q、K、V 全部来自输入xCross‑Attention 则 Q 来自 decoder、K‑V 来自 encoderO(n2)复杂度序列长度变长比如几千 token计算量急剧变大该代码没有加入位置编码如果你要做 GPT/BERT 完整版本需要额外加入Positional Encoding。如果你需要我可以给你写带位置编码残差 LayerNorm完整 Transformer Encoder 单层代码。Q 拿着问题挨个匹配 K匹配分数决定取用 V 的多少通俗拆解Q‑K‑V 逻辑结合数学过程讲明白1. 角色类比Q(Query)提问者每一个 token 带着自己的特征向量提出问题我应该重点关注序列里哪些单词K(Key)候选人的自我介绍序列中每个 token 拿出自身特征用来接受匹配V(Value)候选人真正携带的内容信息。K 只用来打分最终提取内容取自 V。举句子例子He ate appletokenate生成自己的 Q 向量剩下三个 tokenHe、ate、apple各自拿出 KQ 和每个 K 做内积计算相似度Q (ate) 和 K (apple) 内积数值很高Q (ate) 和 K (He) 内积偏低得到原始分数[scoreHe, scoreate, scoreapple]。2. 缩放 Softmax把分数变成权重除以dk缩放防止数值爆炸softmax 归一化所有权重取值 0‑1总和为 1weightssoftmax(dkQK⊤) 这时 w1(He)≈0.1w2(ate)≈0.2w3(apple)≈0.7。含义ate这个单词 70% 的信息来自 apple10% 来自 He。3. 加权 V 获取输出outputatew1⋅VHew2⋅Vatew3⋅Vapple权重越高对应 V 的信息放进来越多权重很低的 V 几乎被忽略。一句话概括你的原话Q 拿着自身特征挨个比对全部 K算出相关性得分softmax 将得分转为权重用权重分配 V 里信息占比这就是自注意力的核心。4. 区分 Self‑Attention 和 Cross‑AttentionSelf‑Attention自注意力Q,K,V全部来自同一组输入句子。句子内部词语互相匹配。Cross‑Attention交叉注意力翻译场景Q 来自 Decoder目标语言K、V 来自 Encoder源语言。例如英译中中文 token 的 Q 去匹配英文单词的 K 和 V。5. 代码层面精简对应关系结合刚才代码attn_score Q K.transpose(-2,-1) / sqrt(d_k) # Q挨个匹配K attn_weight F.softmax(attn_score, dim-1) # 算出权重 out attn_weight V # 根据权重取用V6. 误区纠正Q、K、V 不是原始词向量是输入向量分别乘以可训练矩阵WQ,WK,WV得到模型训练期间自动学习什么样的 Q/K/V 更容易捕捉语义关联内积为什么可以衡量相似度向量夹角越小内积越大特征方向越接近代表语义关联性越强。

**PyTorch 极简实现 Self‑Attention 可运行代码（带注释版）

相关推荐

【企业级GPTs部署白皮书】：从零搭建私有化GPTs工作流，支持SSO集成+审计日志+RBAC权限管控

Gemma手机端Agent落地真相：离线推理与系统级工程挑战

WaveTools鸣潮工具箱：5分钟解锁120帧+画质优化的终极指南

更多资讯

Windows开机加速：5种启动项管理方法详解

Edge兼容性修复实战：从翻车到可复现的完整指南

Aspose.Words for Java 14.9 离线破解版技术解析

企业微信 iPad 方案怎么选？官方 API、自动化能力与风险边界

YOLOv10与红外成像在太阳能板缺陷检测中的应用

PyTorch安装验证全链路指南：CUDA、驱动、Python版本精准校准

今日更新

基于YOLOv11的石油泄漏智能检测系统开发实践

AI行业变局：合规、Mod化与上下文管理的实战启示

计算机视觉与AI原生应用融合的技术实践

本周精选

防幻觉智能体架构：电源硬件设计中的大模型编排与工具化实践

OpenGL面试核心考点：渲染管线、矩阵变换与性能排查全解析

全国计算机二级C语言真题PDF高效刷题指南

本月热门

我发现了一个新思路：用 Remotion + Claude Code 像写代码一样自动化生成短视频

Windows下 Codex 中 Chrome 和 Computer Use 插件不可用问题排查及解决参考方式：TaoToken 统一 Key 配置与验证

2026 大模型集体涨价：用 Python 做企业 Token 成本测算与选型避坑（附配置）

读完文章，想聊聊您的网站？