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

注意力机制

news 2026/7/28 23:59:01

import torch
import torch.nn.functional as F# 一个形状为 (batch_size, seq_len, feature_dim) 的张量（矩阵） x
# 两个样本， 每个样本3 个词语长度，每个词语编码为4维
x = torch.randn(2, 3, 4)  # (batch_size, seq_len, feature_dim)# 定义头数和每个头的维度
num_heads = 2
head_dim = 2# feature_dim 必须是 num_heads * head_dim 的整数倍
assert x.size(-1) == num_heads * head_dim# 定义线性层用于将 x 转换为 Q, K, V 向量
# Wq Wk Wv
linear_q = torch.nn.Linear(4, 4) # 每次创建全链接层，随机参数
linear_k = torch.nn.Linear(4, 4)
linear_v = torch.nn.Linear(4, 4)# 通过线性层计算 Q, K, V
Q = linear_q(x)  # (batch_size, seq_len, feature_dim)
K = linear_k(x)  # (batch_size, seq_len, feature_dim)
V = linear_v(x)  # (batch_size, seq_len, feature_dim)# 将 Q, K, V 分割成 num_heads 个头
def split_heads(tensor, num_heads):# (batch_size, seq_len, feature_dim)  -》 (batch_size, num_heads, seq_len, feature_dim)batch_size, seq_len, feature_dim = tensor.size()head_dim = feature_dim // num_headsoutput = tensor.view(batch_size, seq_len, num_heads, head_dim).transpose(1, 2)# view 维度转换# transpose 维度交换# (batch_size, num_heads, seq_len, feature_dim)return outputQ = split_heads(Q, num_heads)  # (batch_size, num_heads, seq_len, head_dim)
K = split_heads(K, num_heads)  # (batch_size, num_heads, seq_len, head_dim)
V = split_heads(V, num_heads)  # (batch_size, num_heads, seq_len, head_dim)# 计算 Q 和 K 的点积，作为相似度分数 , 也就是自注意力原始权重
# 第i行第j列：第i个词对第j个词的关注程度
raw_weights = torch.matmul(Q, K.transpose(-2, -1))  # (batch_size, num_heads, seq_len, seq_len)# transpose(-2, -1) -》 (batch_size, num_heads, head_dim, seq_len)# 对自注意力原始权重进行缩放， 防止梯度爆炸
scale_factor = K.size(-1) ** 0.5 # 根号 head_dim

scaled_weights = raw_weights / scale_factor  # (batch_size, num_heads, seq_len, seq_len)# 对缩放后的权重进行 softmax 归一化，得到注意力权重
# 归一化操作 / 激活函数
attn_weights = F.softmax(scaled_weights, dim=-1)  # (batch_size, num_heads, seq_len, seq_len)# 将注意力权重应用于 V 向量，计算加权和，得到加权信息
# 每个词的输出 = 该词对所有Value的加权平均
attn_outputs = torch.matmul(attn_weights, V)  # (batch_size, num_heads, seq_len, head_dim)def combine_heads(tensor, num_heads):batch_size, num_heads, seq_len, head_dim = tensor.size()feature_dim = num_heads * head_dimoutput = tensor.transpose(1, 2).contiguous().view(batch_size, seq_len, feature_dim)return output  # (batch_size, seq_len, feature_dim)
attn_outputs = combine_heads(attn_outputs, num_heads)  # (batch_size, seq_len, feature_dim)# 对拼接后的结果进行线性变换
linear_out = torch.nn.Linear(4, 4)
attn_outputs = linear_out(attn_outputs)  # (batch_size, seq_len, feature_dim)
print(" 加权信息 :", attn_outputs)