【pytorch】手把手实现自注意力机制

程思睿Siri

已于 2023-02-25 14:22:40 修改

阅读量3.5k

点赞数 4

CC 4.0 BY-SA版权

分类专栏： python pytorch 模式识别文章标签：深度学习人工智能

于 2022-10-09 18:24:20 首次发布

本文链接：https://blog.csdn.net/m0_53534129/article/details/127229730

python 同时被 3 个专栏收录

8 篇文章

订阅专栏

pytorch

5 篇文章

订阅专栏

模式识别

1 篇文章

订阅专栏

本文介绍了一种基于PyTorch实现的多头自注意力机制方法。该方法接收形状为(B, N, C)的输入矩阵，并通过定义一个多头自注意力类来实现自注意力机制。最终输出经过注意力加权后的特征矩阵。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

背景：

不仅在NLP领域，自注意力机制也在CV领域有着广泛的应用。所以，如何很好地实现自注意力机制成为比较关键的问题。下面我们来对于该机制进行简单实现。

先总结一下思路：

1. 我们的输入是一个（B，N，C）形状的矩阵，其中B代表Batch Size，N代表Time Step，C代表每个Time Step的维度。

2. 我们想做的是，根据输入得到多头的qkv。qkv分别代表query，key，value。我们想用query来查询key而得到一个关联度矩阵A。

3. 由于是多头注意力，我们得到了多个关联度矩阵，我们要将多个关联度矩阵合并为一个。

4. 最后的关联度矩阵和value矩阵相乘，等到最后的输出。

最后的代码如下：

import torch,math
import torch.nn as nn
class MultiHead_SelfAttention(nn.Module):
    def __init__(self, dim, num_head):
        '''

        Args:
            dim: dimension for each time step
            num_head:num head for multi-head self-attention
        '''
        super().__init__()
        self.dim=dim
        self.num_head=num_head
        self.qkv=nn.Linear(dim, dim*3) # extend the dimension for later spliting

    def forward(self, x):
        B, N, C = x.shape
        qkv = self.qkv(x).reshape(B, N, 3, self.num_head, C//self.num_head).permute(2, 0, 3, 1, 4)
        q, k, v= qkv[0], qkv[1], qkv[2]
        att = q@k.transpose(-1, -2)/ math.sqrt(C)
        att = att.softmax(dim=1) # 将多个注意力矩阵合并为一个
        x = (att@v).transpose(1, 2)
        x=x.reshape(B, N, C)
        return x



if __name__=='__main__':
    B = 10
    N = 20
    C = 32
    num_head=8
    x = torch.rand((B, N, C))
    MHSA=Multihead_SelfAttention(C, num_head)
    print(MHSA(x).shape)