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)