QKV加权求和

import torch
import torch.nn.functional as F

#1. 创建 Query、Key 和 Value 张量
q = torch.randn(2, 3, 4) # 形状 (batch_size, seq_len1, feature_dim)
k = torch.randn(2, 4, 4) # 形状 (batch_size, seq_len2, feature_dim)
v = torch.randn(2, 4, 4) # 形状 (batch_size, seq_len2, feature_dim)


# 2. 计算点积，得到原始权重，形状为 (batch_size, seq_len1, seq_len2)
# 批量矩阵乘法 Batched Matrix Multiplication
raw_weights = torch.bmm(q, k.transpose(1, 2))
print("原始权重 :", raw_weights)

# 3. 将原始权重进行缩放（可选），形状仍为 (batch_size, seq_len1, seq_len2)
scaling_factor = q.size(-1) ** 0.5
scaled_weights = raw_weights / scaling_factor
print("缩放后的权重 :", scaled_weights)

# 4. 应用 softmax 函数，使结果的值在 0 和 1 之间，且每一行的和为 1
attn_weights = F.softmax(scaled_weights, dim=-1) # 形状仍为 (batch_size, seq_len1, seq_len2)
print("注意力权重 :", attn_weights)

# 5. 与 Value 相乘，得到注意力分布的加权和 , 形状为 (batch_size, seq_len1, feature_dim)
attn_output = torch.bmm(attn_weights, v)
print("注意力分布的加权和 :", attn_output)