前言

在网上看了一篇文章，借用了文章的大部分代码，并对代码的预测进行修改使得可以正确的预测了，具体链接找了半天找不到

代码

import numpy as np # 导入 numpy 库
import torch # 导入 torch 库
import torch.nn as nn # 导入 torch.nn 库
d_k = 64 # K(=Q) 维度
d_v = 64 # V 维度
# 定义缩放点积注意力类
class ScaledDotProductAttention(nn.Module):def __init__(self):super(ScaledDotProductAttention, self).__init__()        def forward(self, Q, K, V, attn_mask):#------------------------- 维度信息 --------------------------------        # Q K V [batch_size, n_heads, len_q/k/v, dim_q=k/v] (dim_q=dim_k)# attn_mask [batch_size, n_heads, len_q, len_k]#----------------------------------------------------------------# 计算注意力分数（原始权重）[batch_size，n_heads，len_q，len_k]scores = torch.matmul(Q, K.transpose(-1, -2)) / np.sqrt(d_k) #------------------------- 维度信息 --------------------------------        # scores [batch_size, n_heads, len_q, len_k]#-----------------------------------------------------------------        # 使用注意力掩码，将 attn_mask 中值为 1 的位置的权重替换为极小值#------------------------- 维度信息 -------------------------------- # attn_mask [batch_size, n_heads, len_q, len_k], 形状和 scores 相同#-----------------------------------------------------------------    scores.masked_fill_(attn_mask, -1e9) # 对注意力分数进行 softmax 归一化weights = nn.Softmax(dim=-1)(scores) #------------------------- 维度信息 -------------------------------- # weights [batch_size, n_heads, len_q, len_k], 形状和 scores 相同#-----------------------------------------------------------------         # 计算上下文向量（也就是注意力的输出）, 是上下文信息的紧凑表示context = torch.matmul(weights, V) #------------------------- 维度信息 -------------------------------- # context [batch_size, n_heads, len_q, dim_v]#-----------------------------------------------------------------    return context, weights # 返回上下文向量和注意力分数# 定义多头自注意力类
d_embedding = 512  # Embedding 的维度
n_heads = 8  # Multi-Head Attention 中头的个数
batch_size = 6 # 每一批的数据大小
class MultiHeadAttention(nn.Module):def __init__(self):super(MultiHeadAttention, self).__init__()self.W_Q = nn.Linear(d_embedding, d_k * n_heads) # Q的线性变换层self.W_K = nn.Linear(d_embedding, d_k * n_heads) # K的线性变换层self.W_V = nn.Linear(d_embedding, d_v * n_heads) # V的线性变换层self.linear = nn.Linear(n_heads * d_v, d_embedding)self.layer_norm = nn.LayerNorm(d_embedding)def forward(self, Q, K, V, attn_mask): #------------------------- 维度信息 -------------------------------- # Q K V [batch_size, len_q/k/v, embedding_dim] #-----------------------------------------------------------------        residual, batch_size = Q, Q.size(0) # 保留残差连接# 将输入进行线性变换和重塑，以便后续处理q_s = self.W_Q(Q).view(batch_size, -1, n_heads, d_k).transpose(1,2)        k_s = self.W_K(K).view(batch_size, -1, n_heads, d_k).transpose(1,2)v_s = self.W_V(V).view(batch_size, -1, n_heads, d_v).transpose(1,2)#------------------------- 维度信息 -------------------------------- # q_s k_s v_s: [batch_size, n_heads, len_q/k/v, d_q=k/v]#----------------------------------------------------------------- # 将注意力掩码复制到多头 attn_mask: [batch_size, n_heads, len_q, len_k]attn_mask = attn_mask.unsqueeze(1).repeat(1, n_heads, 1, 1<