注意力机制练习 Attention

 

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_heads
    output = tensor.view(batch_size, seq_len, num_heads, head_dim).transpose(1, 2)
    # view 维度转换
    # transpose 维度交换
    # (batch_size, num_heads, seq_len, feature_dim)
    return output


Q = 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_dim
    output = 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)