transformer代码

import torch
import numpy as np

#定义输入数据
x = [
  [1, 0, 1, 0], # Input 1
  [0, 2, 0, 2], # Input 2
  [1, 1, 1, 1]  # Input 3
 ]
x = torch.tensor(np.array(x),dtype=torch.float32)

#初始化KQV的权重
w_key = [
  [0, 0, 1],
  [1, 1, 0],
  [0, 1, 0],
  [1, 1, 0]
]
w_query = [
  [1, 0, 1],
  [1, 0, 0],
  [0, 0, 1],
  [0, 1, 1]
]
w_value = [
  [0, 2, 0],
  [0, 3, 0],
  [1, 0, 3],
  [1, 1, 0]
]
w_key = torch.tensor(np.array(w_key), dtype=torch.float32)
w_query = torch.tensor(np.array(w_query), dtype=torch.float32)
w_value = torch.tensor(np.array(w_value), dtype=torch.float32)

#推导键、查询和值
keys = x @ w_key
querys = x @ w_query
values = x @ w_value

#计算注意力得分
attn_scores = querys @ keys.T

#对得分做归一化
from torch.nn.functional import softmax
attn_scores = softmax(attn_scores,dim=-1)

#将得分和值进行相乘
weighted_values = values[:,None] * attn_scores.T[:,:,None]

#对加权值进行求和
outputs = weighted_values.sum(dim=0)
print(outputs)