手撕SelfAttention

最新推荐文章于 2024-11-17 12:01:11 发布

云帆@

最新推荐文章于 2024-11-17 12:01:11 发布

阅读量515

点赞数 2

分类专栏： torch 文章标签：深度学习 pytorch 人工智能

本文链接：https://blog.csdn.net/weixin_40777649/article/details/136249890

版权

torch 专栏收录该内容

21 篇文章 1 订阅

订阅专栏

本文详细介绍了如何使用PyTorch实现自注意力机制（Self-Attention），包括定义模块、线性变换以及softmax计算注意力权重。作者通过实例展示了如何构造并运行一个自注意力层，适用于深度学习中的序列建模任务。

摘要由CSDN通过智能技术生成

一、目录

定义
代码实现

二、实现

定义
代码实现


#手撕 self attention
import torch
import torch.nn as nn
import numpy as np

class SelfAttention(nn.Module):
    def __init__(self,hidden_dim,dim_q,dim_v):
        super(SelfAttention,self).__init__()
        self.hidden_dim=hidden_dim
        self.dim_q=dim_q
        self.dim_k=dim_q
        self.dim_v=dim_v

        self.linear_q=nn.Linear(self.hidden_dim,self.dim_q)
        self.linear_k=nn.Linear(self.hidden_dim,self.dim_k)
        self.linear_v=nn.Linear(self.hidden_dim,self.dim_v)
        self.norm_fact=1/np.sqrt(self.dim_k)                #保持均值、方差不变，使得训练过程中梯度值保持稳定

    def forward(self,x):

        q=self.linear_q(x)                                  #为了提升模型的拟合能力，矩阵W都是可以训练的，起到一个缓冲的效果。
        k=self.linear_k(x)
        v=self.linear_v(x)

        acore=torch.matmul(q,k.transpose(1,2))*self.norm_fact
                                              #内积：以行向量的角度理解，里面保存了每个向量与自己和其他向量进行内积运算的结果，代表词的相关性
        a=torch.softmax(acore,dim=-1)
        att=torch.matmul(a,v)
        return att

if __name__ == '__main__':
    batch=2
    seq_len=5
    hidden_dim=4
    x=torch.randn(batch,seq_len,hidden_dim)
    attention=SelfAttention(hidden_dim,10,hidden_dim)
    print(attention(x).shape)