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意思是一样的:

哪里没理解到位?

相关推荐
old_power3 分钟前
PyTorch 分布式训练(Distributed Data Parallel, DDP)简介
pytorch·ddp
Sapphire~4 分钟前
odoo-045 ModuleNotFoundError: No module named ‘_sqlite3‘
python·ubuntu·odoo
Y1nhl1 小时前
搜广推校招面经六十四
人工智能·深度学习·leetcode·广告算法·推荐算法·搜索算法
Y1nhl2 小时前
Pyspark学习一:概述
数据库·人工智能·深度学习·学习·spark·pyspark·大数据技术
维度攻城狮4 小时前
实现在Unity3D中仿真汽车,而且还能使用ros2控制
python·unity·docker·汽车·ros2·rviz2
简简单单做算法4 小时前
基于mediapipe深度学习和限定半径最近邻分类树算法的人体摔倒检测系统python源码
人工智能·python·深度学习·算法·分类·mediapipe·限定半径最近邻分类树
hvinsion5 小时前
基于PyQt5的自动化任务管理软件:高效、智能的任务调度与执行管理
开发语言·python·自动化·自动化任务管理
就决定是你啦!5 小时前
机器学习 第一章 绪论
人工智能·深度学习·机器学习
飞飞翼6 小时前
python-flask
后端·python·flask
林九生7 小时前
【Python】Browser-Use:让 AI 替你掌控浏览器,开启智能自动化新时代!
人工智能·python·自动化