当前位置：首页 > news >正文

桂林网站制作网站wordpress主题noren

news 2025/9/27 16:40:09

桂林网站制作网站,wordpress主题noren,增城专业建站公司,wordpress 5.0.4自注意力（Self-Attention）与交叉注意力（Cross-Attention）PyTorch 简单实现在深度学习中，注意力机制是现代 Transformer 架构的核心思想之一。本文将介绍两种常见的注意力机制：自注意力（Self-A…

自注意力（Self-Attention）与交叉注意力（Cross-Attention）PyTorch 简单实现

在深度学习中，注意力机制是现代 Transformer 架构的核心思想之一。本文将介绍两种常见的注意力机制：自注意力（Self-Attention）与交叉注意力（Cross-Attention），并通过 PyTorch 给出简单实现与使用示例。

📦 必要导入

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import einsum
from einops import rearrange, repeat
from inspect import isfunction# 一些基础工具函数
def exists(val):return val is not Nonedef default(val, d):if exists(val):return valreturn d() if isfunction(d) else d

🔍 什么是注意力机制？

注意力机制允许模型在处理输入序列时自动聚焦于最相关的部分，从而增强建模能力。以 Transformer 为例，它通过注意力机制建立了序列中不同位置之间的信息关联。

🤖 自注意力（Self-Attention）

自注意力是指查询（Query）、键（Key）、值（Value）都来自同一个输入序列。这种机制允许序列中的每个元素关注其它所有位置的信息，是 BERT、GPT 等模型的基本构件。

✅ PyTorch 实现：

class SelfAttention(nn.Module):def __init__(self, dim, heads=8, dim_head=64, dropout=0.0):super().__init__()inner_dim = dim_head * headsself.scale = dim_head ** -0.5self.heads = headsself.to_qkv = nn.Linear(dim, inner_dim * 3, bias=False)self.to_out = nn.Sequential(nn.Linear(inner_dim, dim),nn.Dropout(dropout))def forward(self, x, mask=None):h = self.headsqkv = self.to_qkv(x)q, k, v = qkv.chunk(3, dim=-1)q, k, v = map(lambda t: rearrange(t, "b n (h d) -> (b h) n d", h=h), (q, k, v))sim = einsum("b i d, b j d -> b i j", q, k) * self.scaleif exists(mask):...attn = sim.softmax(dim=-1)out = einsum("b i j, b j d -> b i d", attn, v)out = rearrange(out, "(b h) n d -> b n (h d)", h=h)return self.to_out(out)

🧪 使用示例：

attn = SelfAttention(dim=64)
x = torch.randn(1, 10, 64)
out = attn(x)
print(out.shape)  # torch.Size([1, 10, 64])

🔁 交叉注意力（Cross-Attention）

交叉注意力允许模型在处理一个序列时，从另一个序列中获取信息。常用于：

编码器-解码器结构（如 Transformer 翻译模型）
图文跨模态对齐
条件生成任务

与自注意力的不同在于：

Query 来自当前输入（例如解码器）
Key 与 Value 来自另一个序列（例如编码器）

✅ PyTorch 实现：

class CrossAttention(nn.Module):def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.0):super().__init__()inner_dim = dim_head * headscontext_dim = default(context_dim, query_dim)self.scale = dim_head ** -0.5self.heads = headsself.to_q = nn.Linear(query_dim, inner_dim, bias=False)self.to_k = nn.Linear(context_dim, inner_dim, bias=False)self.to_v = nn.Linear(context_dim, inner_dim, bias=False)self.to_out = nn.Sequential(nn.Linear(inner_dim, query_dim),nn.Dropout(dropout))def forward(self, x, context=None, mask=None):h = self.headscontext = default(context, x)q = self.to_q(x)k = self.to_k(context)v = self.to_v(context)q, k, v = map(lambda t: rearrange(t, "b n (h d) -> (b h) n d", h=h), (q, k, v))sim = einsum("b i d, b j d -> b i j", q, k) * self.scaleif exists(mask):...attn = sim.softmax(dim=-1)out = einsum("b i j, b j d -> b i d", attn, v)out = rearrange(out, "(b h) n d -> b n (h d)", h=h)return self.to_out(out)

🧪 使用示例：

ca = CrossAttention(query_dim=64, context_dim=77)
x = torch.randn(1, 10, 64)          # 解码器输入
context = torch.randn(1, 20, 77)    # 编码器输出
out = ca(x, context)
print(out.shape)  # torch.Size([1, 10, 64])

🧠 总结对比

模块	Query 来自	Key/Value 来自	典型应用
Self-Attention	当前输入	当前输入	BERT、GPT、自注意图像建模
Cross-Attention	当前输入	外部上下文	编解码结构、跨模态、条件生成