python手撕代码——多头注意力机制

备注：以前搞的方法比较传统，现在搞卫星图像目标跟踪要用transfomer，最近在学习这方面，为了加深理解，打算手写一个transfomrer，发帖子当作笔记使用，有错误还请指出。

参考帖子1.
参考帖子2
b站还有大佬的视频手撕讲解

1.定义变量，以及main函数调用多头注意力类

假设输入10个单词，我们把他embedding为输入变量X，这个X为max_len乘embedding_dim 的矩阵

embedding_dim = 512                              #每个单词的维度
max_len = 10                                                 #单词数量
num_heads = 8                                             #多头注意力机制的头数
embedded_sentence=X                           #这就是输入，10x512的矩阵

接下来我们定义一个多头注意力的类，main函数如下

attention_layer = MultiHeadAttention(embed_dim=embedding_dim, num_heads=num_heads)
attention_output = attention_layer(embedded_sentence.unsqueeze(0))  # 添加 batch 维度

2.多头注意力类的实现

2.1初始化函数

def __init__(self, embed_dim, num_heads):
        super(MultiHeadAttention, self).__init__()
        self.num_heads = num_heads
        self.embed_dim = embed_dim
        self.head_dim = embed_dim // num_heads     #这一步将输入向量分割为多个注意力头，每个头独立地执行注意力计算
        

        self.query=nn.Linear(embed_dim, embed_dim)
        self.key=nn.Linear(embed_dim, embed_dim)
        self.value=nn.Linear(embed_dim, embed_dim)
        self.out = nn.Linear(embed_dim, embed_dim)

    #nn.Linear(embed_dim, embed_dim)定义一些输入和输出维度都为embed_dim的特征映射，这些特征映射都是设置一个随机初始值，然后通过训练得到。 

2.2 forward函数

2.2.1线性映射

懒得画图了，借用一下参考帖子1里的图。。。。。

def forward(self,X):
        batch_size, seq_len, embed_dim = X.size()
         #线性映射
        q = self.query(X)
        k = self.key(X)
        v = self.value(X)
        print("q1 shape:", q.shape)
        print("k1 shape:", k.shape)

通过线性映射得到qkv，也就是X*$w^q$=q 的过程

这里输出：
q1 shape: torch.Size([1, 10, 512])
k1 shape: torch.Size([1, 10, 512])

2.2.2 多头运算

图是借用 参考帖子1 的。。。。。。。。。。。。。。。。

接下来分为num_heads个头计算注意力，也就是将矩阵进行变形。
[batch_size, seq_len, embed_dim]维度变为[batch_size, seq_len, num_heads, head_dim]
transpose(1, 2) 调换了 seq_len 和 num_heads 的维度[batch_size, num_heads, seq_len, head_dim]

        q=q.view(batch_size, seq_len