DL之self-attention：self-attention自注意力机制模块思路的8个步骤及其代码实现

1. import torch
2. 
3. #1、准备输入：Input 1、2、3
4. x = [[1, 0, 1, 0], 
5.      [0, 2, 0, 2],
6.      [1, 1, 1, 1]  ]
7. x = torch.tensor(x, dtype=torch.float32)
8. 
9. 
10. #2、初始化权重
11. w_key   = [ [0, 0, 1], [1, 1, 0], [0, 1, 0], [1, 1, 0] ]
12. w_query = [ [1, 0, 1], [1, 0, 0], [0, 0, 1], [0, 1, 1] ]
13. w_value = [ [0, 2, 0], [0, 3, 0], [1, 0, 3], [1, 1, 0] ]
14. 
15. w_key   = torch.tensor(w_key,   dtype=torch.float32)
16. w_query = torch.tensor(w_query, dtype=torch.float32)
17. w_value = torch.tensor(w_value, dtype=torch.float32)
18. 
19. 
20. #3、推导键、查询和值
21. keys   = x @ w_key
22. querys = x @ w_query
23. values = x @ w_value
24. 
25. print(keys)   # tensor([[0., 1., 1.], [4., 4., 0.], [2., 3., 1.]])
26. print(querys) # tensor([[1., 0., 2.], [2., 2., 2.], [2., 1., 3.]])
27. print(values) # tensor([[1., 2., 3.], [2., 8., 0.], [2., 6., 3.]])
28. 
29. 
30. #4、计算注意力得分
31. attn_scores = querys @ keys.t()
32. # tensor([[ 2.,  4.,  4.],  # attention scores from Query 1
33. #         [ 4., 16., 12.],  # attention scores from Query 2
34. #         [ 4., 12., 10.]]) # attention scores from Query 3
35. 
36. 
37. 
38. #5、计算softmax
39. from torch.nn.functional import softmax
40. attn_scores_softmax = softmax(attn_scores, dim=-1)
41. print('attn_scores_softmax：','\n',attn_scores_softmax)
42. # tensor([[6.3379e-02, 4.6831e-01, 4.6831e-01],
43. #         [6.0337e-06, 9.8201e-01, 1.7986e-02],
44. #         [2.9539e-04, 8.8054e-01, 1.1917e-01]])
45. 
46. # For readability, approximate the above as follows
47. attn_scores_softmax = [  [0.0, 0.5, 0.5], [0.0, 1.0, 0.0], [0.0, 0.9, 0.1] ]
48. attn_scores_softmax = torch.tensor(attn_scores_softmax)
49. 
50. #6、将得分和值相乘
51. weighted_values = values[:,None] * attn_scores_softmax.t()[:,:,None]
52. print('weighted_values：','\n',weighted_values)
53. # tensor([[[0.0000, 0.0000, 0.0000], [0.0000, 0.0000, 0.0000], [0.0000, 0.0000, 0.0000]],
54. #         [[1.0000, 4.0000, 0.0000], [2.0000, 8.0000, 0.0000], [1.8000, 7.2000, 0.0000]],
55. #         [[1.0000, 3.0000, 1.5000], [0.0000, 0.0000, 0.0000], [0.2000, 0.6000, 0.3000]]])
56. 
57. #7、求和加权值
58. outputs = weighted_values.sum(dim=0)
59. # tensor([[2.0000, 7.0000, 1.5000], [2.0000, 8.0000, 0.0000], [2.0000, 7.8000, 0.3000]]) # Output1、2、3
60. print('outputs：','\n',outputs)
61. 
62. 
63. 
64.