SelfAttention和MultiHeadAttion实现demo

#encoding:utf-8

from math import sqrt

import torch

import torch.nn as nn

class Self_Attention(nn.Module):

def init(self, input_dim, dim_k, dim_v):

super(Self_Attention, self). init()

self.q = nn.Linear(input_dim, dim_k)

self.k = nn.Linear(input_dim, dim_k)

self.v = nn.Linear(input_dim, dim_v)

self.norm_fact = 1 / sqrt(dim_k)

def forward(self, x):

print("x.shape:", x.shape)

print("q.shape:", self.q.shape)

Q = self.q(x)

print("Q.shape:", Q.shape)

K = self.k(x)

print("K.shape:", K.shape)

V = self.v(x)

print("V.shape:", V.shape)

atten = nn.Softmax(dim=-1)(torch.bmm(Q,K.permute(0,2,1))) * self.norm_fact

output = torch.bmm(atten, V)

return output

print("\n")

print("self attention:")

x = torch.randn(4,3,1024)

print(x)

print("input size:", x.size())

self_attention = Self_Attention(1024,128,5)

res = self_attention(x)

print("\n")

print(res)

print("output size:", res.size())

print("\n")

class Self_Attention_Muti_Head(nn.Module):

def init(self, input_dim, dim_k, dim_v, nums_head):

super(Self_Attention_Muti_Head, self).init()

assert dim_k % nums_head == 0

assert dim_v % nums_head == 0

self.q = nn.Linear(input_dim, dim_k)

self.k = nn.Linear(input_dim, dim_k)

self.v = nn.Linear(input_dim, dim_v)

self.nums_head = nums_head

self.dim_k = dim_k

self.dim_v = dim_v

self._norm_fact = 1 / sqrt(dim_k)

def forward(self, x):

Q = self.q(x).reshape(-1, x.shape[0], x.shape[1], self.dim_k//self.nums_head)

K = self.k(x).reshape(-1, x.shape[0], x.shape[1], self.dim_k//self.nums_head)

V = self.v(x).reshape(-1, x.shape[0], x.shape[1], self.dim_v//self.nums_head)

print("x.shape:", x.shape)

print("Q.shape", Q.size())

atten = nn.Softmax(dim=-1)(torch.matmul(Q, K.permute(0,1,3,2)))

output = torch.matmul(atten, V).reshape(x.shape[0], x.shape[1], -1)

return output

print("\n")

print("multi head attention:")

x = torch.randn(4,3,1024)

print(x)

print(x.size())

self_attention = Self_Attention_Muti_Head(1024,128,6,2)

res = self_attention(x)

print("\n")

print(res)

print(res.size())


有个问题:

根据文献:https://arxiv.org/pdf/1911.02150.pdf,感觉这里说的Multi Head Attenion和 Group Query Attention意思是一样的:

这下面这张经典的图中的的Grouped-query意思是一样的:

哪里没理解到位?

相关推荐
noravinsc2 分钟前
django 一个表中包括id和parentid,如何通过parentid找到全部父爷id
python·django·sqlite
ajassi200013 分钟前
开源 python 应用 开发(三)python语法介绍
linux·python·开源·自动化
沉默媛35 分钟前
如何安装python以及jupyter notebook
开发语言·python·jupyter
HollowKnightZ37 分钟前
目标姿态估计综述:Deep Learning-Based Object Pose Estimation: A Comprehensive Survey
人工智能·深度学习
加油吧zkf1 小时前
Conda虚拟环境管理:从入门到精通的常用命令
图像处理·深度学习·计算机视觉·conda
Deng9452013142 小时前
基于Python的旅游数据可视化应用
python·numpy·pandas·旅游·数据可视化技术
2401_878624792 小时前
pytorch 自动微分
人工智能·pytorch·python·机器学习
胖达不服输2 小时前
「日拱一码」021 机器学习——特征工程
人工智能·python·机器学习·特征工程
小哥谈3 小时前
论文解析篇 | YOLOv12:以注意力机制为核心的实时目标检测算法
人工智能·深度学习·yolo·目标检测·机器学习·计算机视觉
screenCui3 小时前
macOS运行python程序遇libiomp5.dylib库冲突错误解决方案
开发语言·python·macos