用einsum实现MultiHeadAttention前向传播

2024-09-08 14:04
文章标签 实现 传播 前向 einsum multiheadattention

本文主要是介绍用einsum实现MultiHeadAttention前向传播，希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！

einsum教程网站Einstein Summation in Numpy | Olexa Bilaniuk's IFT6266H16 Course Blog

编写训练模型

import tensorflow as tfclass Model(tf.keras.Model):def __init__(self, num_heads, model_dim):super().__init__()self.AttentionLayer = tf.keras.layers.MultiHeadAttention(num_heads=num_heads, key_dim=model_dim)self.OutputLayer = tf.keras.layers.Dense(units=1)def call(self, x):x = self.AttentionLayer(query=x, value=x)x = self.OutputLayer(x)return xmodel = Model(num_heads=2, model_dim=4)model.compile(loss=tf.keras.losses.BinaryCrossentropy(from_logits=True),optimizer="Adam",metrics=['accuracy'])input_train = tf.constant([[[1, 2, 3], [4, 5, 6]],[[1, 1, 1], [2, 2, 6]]], dtype=tf.float32)output_label = tf.constant([[[1], [0]],[[0], [0]]], dtype=tf.float32)tf_callback = tf.keras.callbacks.TensorBoard(log_dir="./logs")
model.fit(x=input_train,y=output_label,epochs=10,callbacks=[tf_callback])tf.saved_model.save(model, 'MultiHeadAttention')

训练样本input_train的shape为(2, 2, 3)→(batch_size, sentence_length, Embedding_dim)，经过模型后的输出shape为(2,2,1)，标签值的shape为(2, 2, 1)，损失函数选择的是二分类交叉熵。所以该模型可以应用于二分类性质的命名实体识别。简单来说，该模型可以判断一个句子中哪些词属于我们感兴趣的词类，哪些不属于。

打印模型矩阵参数

import tensorflow as tfsave_path = 'MultiHeadAttention/variables/variables'  #reader = tf.train.load_checkpoint(save_path)  # 得到CheckpointReader"""  打印Checkpoint中存储的所有参数名和参数shape """
for variable_name, variable_shape in reader.get_variable_to_shape_map().items():print(f'{variable_name} : {variable_shape}')打印结果，_CHECKPOINTABLE_OBJECT_GRAPH : []
optimizer/beta_1/.ATTRIBUTES/VARIABLE_VALUE : []
keras_api/metrics/0/count/.ATTRIBUTES/VARIABLE_VALUE : []
variables/5/.OPTIMIZER_SLOT/optimizer/v/.ATTRIBUTES/VARIABLE_VALUE : [2, 4]
optimizer/beta_2/.ATTRIBUTES/VARIABLE_VALUE : []
optimizer/learning_rate/.ATTRIBUTES/VARIABLE_VALUE : []
keras_api/metrics/0/total/.ATTRIBUTES/VARIABLE_VALUE : []
keras_api/metrics/1/count/.ATTRIBUTES/VARIABLE_VALUE : []
keras_api/metrics/1/total/.ATTRIBUTES/VARIABLE_VALUE : []
optimizer/decay/.ATTRIBUTES/VARIABLE_VALUE : []
variables/2/.OPTIMIZER_SLOT/optimizer/v/.ATTRIBUTES/VARIABLE_VALUE : [3, 2, 4]
optimizer/iter/.ATTRIBUTES/VARIABLE_VALUE : []
variables/0/.ATTRIBUTES/VARIABLE_VALUE : [3, 2, 4]
variables/0/.OPTIMIZER_SLOT/optimizer/m/.ATTRIBUTES/VARIABLE_VALUE : [3, 2, 4]
variables/0/.OPTIMIZER_SLOT/optimizer/v/.ATTRIBUTES/VARIABLE_VALUE : [3, 2, 4]
variables/1/.ATTRIBUTES/VARIABLE_VALUE : [2, 4]
variables/1/.OPTIMIZER_SLOT/optimizer/m/.ATTRIBUTES/VARIABLE_VALUE : [2, 4]
variables/5/.ATTRIBUTES/VARIABLE_VALUE : [2, 4]
variables/1/.OPTIMIZER_SLOT/optimizer/v/.ATTRIBUTES/VARIABLE_VALUE : [2, 4]
variables/6/.OPTIMIZER_SLOT/optimizer/v/.ATTRIBUTES/VARIABLE_VALUE : [2, 4, 3]
variables/2/.ATTRIBUTES/VARIABLE_VALUE : [3, 2, 4]
variables/2/.OPTIMIZER_SLOT/optimizer/m/.ATTRIBUTES/VARIABLE_VALUE : [3, 2, 4]
variables/3/.ATTRIBUTES/VARIABLE_VALUE : [2, 4]
variables/3/.OPTIMIZER_SLOT/optimizer/m/.ATTRIBUTES/VARIABLE_VALUE : [2, 4]
variables/3/.OPTIMIZER_SLOT/optimizer/v/.ATTRIBUTES/VARIABLE_VALUE : [2, 4]
variables/4/.ATTRIBUTES/VARIABLE_VALUE : [3, 2, 4]
variables/4/.OPTIMIZER_SLOT/optimizer/m/.ATTRIBUTES/VARIABLE_VALUE : [3, 2, 4]
variables/4/.OPTIMIZER_SLOT/optimizer/v/.ATTRIBUTES/VARIABLE_VALUE : [3, 2, 4]
variables/5/.OPTIMIZER_SLOT/optimizer/m/.ATTRIBUTES/VARIABLE_VALUE : [2, 4]
variables/9/.OPTIMIZER_SLOT/optimizer/m/.ATTRIBUTES/VARIABLE_VALUE : [1]
variables/6/.ATTRIBUTES/VARIABLE_VALUE : [2, 4, 3]
variables/6/.OPTIMIZER_SLOT/optimizer/m/.ATTRIBUTES/VARIABLE_VALUE : [2, 4, 3]
variables/7/.ATTRIBUTES/VARIABLE_VALUE : [3]
variables/8/.ATTRIBUTES/VARIABLE_VALUE : [3, 1]
variables/7/.OPTIMIZER_SLOT/optimizer/m/.ATTRIBUTES/VARIABLE_VALUE : [3]
variables/7/.OPTIMIZER_SLOT/optimizer/v/.ATTRIBUTES/VARIABLE_VALUE : [3]
variables/8/.OPTIMIZER_SLOT/optimizer/m/.ATTRIBUTES/VARIABLE_VALUE : [3, 1]
variables/8/.OPTIMIZER_SLOT/optimizer/v/.ATTRIBUTES/VARIABLE_VALUE : [3, 1]
variables/9/.ATTRIBUTES/VARIABLE_VALUE : [1]
variables/9/.OPTIMIZER_SLOT/optimizer/v/.ATTRIBUTES/VARIABLE_VALUE : [1]

其中我们只需要注意如下参数，

variables/0/.ATTRIBUTES/VARIABLE_VALUE : [3, 2, 4]

variables/1/.ATTRIBUTES/VARIABLE_VALUE : [2, 4]

variables/5/.ATTRIBUTES/VARIABLE_VALUE : [2, 4]

variables/2/.ATTRIBUTES/VARIABLE_VALUE : [3, 2, 4]

variables/3/.ATTRIBUTES/VARIABLE_VALUE : [2, 4]

variables/4/.ATTRIBUTES/VARIABLE_VALUE : [3, 2, 4]

variables/6/.ATTRIBUTES/VARIABLE_VALUE : [2, 4, 3]

variables/7/.ATTRIBUTES/VARIABLE_VALUE : [3]
variables/8/.ATTRIBUTES/VARIABLE_VALUE : [3, 1]

variables/9/.ATTRIBUTES/VARIABLE_VALUE : [1]

这些参数代表着模型中所有的kernel和bias，variable后面的序号代表着输入在模型中依次经历的运算。利用模型内部的计算逻辑以及shape分布，我们可以找出对应的kernel和bias。

值得注意的是，在tf.keras.layers.MultiHeadAttention源码中，query映射层是先进行计算的，后面依次为key映射层，value映射层。所以，

variables/0/.ATTRIBUTES/VARIABLE_VALUE : [3, 2, 4] 对应query映射层的kernel

variables/1/.ATTRIBUTES/VARIABLE_VALUE : [2, 4]对应query映射层的bias

variables/2/.ATTRIBUTES/VARIABLE_VALUE : [3, 2, 4]对应key映射层的kernel

variables/3/.ATTRIBUTES/VARIABLE_VALUE : [2, 4]对应key映射层的bias

variables/4/.ATTRIBUTES/VARIABLE_VALUE : [3, 2, 4]对应value映射层的kernel

variables/5/.ATTRIBUTES/VARIABLE_VALUE : [2, 4]对应value映射层的bias

variables/6/.ATTRIBUTES/VARIABLE_VALUE : [2, 4, 3]对应MutiHeadAttention中输出映射层的kernel

variables/7/.ATTRIBUTES/VARIABLE_VALUE : [3]对应MutiHeadAttention中输出映射层的bias

variables/8/.ATTRIBUTES/VARIABLE_VALUE : [3, 1]对应Dense层的kernel

variables/9/.ATTRIBUTES/VARIABLE_VALUE : [1]对应Dense层的bias

MultiHeadAttention

多头映射层

''' abc , cde -> abde a : size of batchb : length of sequencec : dimension of Embeddingd : number of headse : dimension of output x : input, shape=(batch_size, sequence_length, dim)y : kenel, shape=(dim, heads_num, output_dim) '''def project_dense(x, y):output = np.empty((x.shape[0], x.shape[1], y.shape[1], y.shape[2]))for a in range(x.shape[0]):for b in range(x.shape[1]):for d in range(y.shape[1]):for e in range(y.shape[2]):sum_dot = 0for c in range(x.shape[2]):sum_dot += x[a][b][c] * y[c][d][e]output[a][b][d][e] = sum_dotreturn output

代码验证，

''' shape = (2, 2, 3) '''
x = tf.constant([[[1, 2, 3], [4, 5, 3]],[[7, 8, 9], [10, 11, 12]]])''' shape = (3, 2, 3) '''
y = tf.constant([[[1, 2, 3], [4, 5, 3]],[[7, 8, 9], [10, 11, 12]],[[7, 8, 9], [10, 11, 12]]])print(tf.einsum('abc,cde->abde', x, y).numpy())
print(project_dense(x, y))[[[[ 36  42  48][ 54  60  63]][[ 60  72  84][ 96 108 108]]][[[126 150 174][198 222 225]][[171 204 237][270 303 306]]]][[[[ 36.  42.  48.][ 54.  60.  63.]][[ 60.  72.  84.][ 96. 108. 108.]]][[[126. 150. 174.][198. 222. 225.]][[171. 204. 237.][270. 303. 306.]]]]

多头注意力层

计算注意力分数矩阵
''' aecd, abcd -> acbe a : size of batch e : number of key b : number of queryc : number of heads d : dimension of query/key x is key and y is query '''def compute_AttentionScores(x, y):output = np.empty((x.shape[0], x.shape[2], y.shape[1], x.shape[1]))for a in range(x.shape[0]):for c in range(x.shape[2]):for b in range(y.shape[1]):for e in range(x.shape[1]):sum_dot = 0for d in range(x.shape[3]):sum_dot += x[a][e][c][d] * y[a][b][c][d]output[a][c][b][e] = sum_dotreturn output
根据注意力分数对value进行加权叠加
''' acbe,aecd->abcd a : size of batch c : number of headsb : number of querye : number of key/value d : dimension of value x is attention_scores and y is value'''def Value_WeightedStack(x, y):output = np.empty((x.shape[0], x.shape[2], y.shape[2], y.shape[3]))for a in range(x.shape[0]):for b in range(x.shape[2]):for c in range(y.shape[2]):for d in range(y.shape[3]):sum_dot = 0for e in range(x.shape[3]):sum_dot += x[a][c][b][e] * y[a][e][c][d]output[a][b][c][d] = sum_dotreturn output

输出映射层

''' abcd, cde -> abe a : size of batch b : number of queryc : number of head d : dimension of value e : dimension of output x is WeightedStack_Value and y is kernel'''def project_final(x, y):output = np.empty((x.shape[0], x.shape[1], y.shape[2]))for a in range(x.shape[0]):for b in range(x.shape[1]):for e in range(y.shape[2]):sum_dot = 0for c in range(x.shape[2]):for d in range(x.shape[3]):sum_dot += x[a][b][c][d] * y[c][d][e]output[a][b][e] = sum_dotreturn output

Dense

为了使得方便验证这次试验，在MultiHeadAttentionn后面添加一个Dense层。

'''abc, cd -> abd x is input and y is kernel '''def output_dense(x, y):output = np.empty((x.shape[0], x.shape[1], y.shape[1]))for a in range(x.shape[0]):for b in range(x.shape[1]):for d in range(y.shape[1]):sum_dot = 0for c in range(x.shape[2]):sum_dot += x[a][b][c] * y[c][d]output[a][b][d] = sum_dotreturn output

构建模型的前向传播

在前面我们找出了模型内部所有的kenel和bias，接下来我们将打印出这些参数，并将这些参数添加到我们自己编写的模型中去，实现前向传播。

print(reader.get_tensor('variables/0/.ATTRIBUTES/VARIABLE_VALUE')) //k1
print(reader.get_tensor('variables/1/.ATTRIBUTES/VARIABLE_VALUE')) //b1
print(reader.get_tensor('variables/5/.ATTRIBUTES/VARIABLE_VALUE')) //b2
print(reader.get_tensor('variables/2/.ATTRIBUTES/VARIABLE_VALUE')) //k2
print(reader.get_tensor('variables/3/.ATTRIBUTES/VARIABLE_VALUE')) //b3
print(reader.get_tensor('variables/4/.ATTRIBUTES/VARIABLE_VALUE')) //k3
print(reader.get_tensor('variables/6/.ATTRIBUTES/VARIABLE_VALUE')) //MultiHead_output_kerner
print(reader.get_tensor('variables/7/.ATTRIBUTES/VARIABLE_VALUE')) //MultiHead_output_bias
print(reader.get_tensor('variables/8/.ATTRIBUTES/VARIABLE_VALUE')) //output_dense_kernel
print(reader.get_tensor('variables/9/.ATTRIBUTES/VARIABLE_VALUE')) //output_dense_bias
class my_model:def __init__(self, input):self.input = inputdef __call__(self):x = tf.cast(self.input,dtype=tf.double)value = tf.add(project_dense(x, k3), b2)key = tf.add(project_dense(x, k2), b3)query = tf.add(project_dense(x, k1), b1)attention_scores = tf.nn.softmax(compute_AttentionScores(key, query), axis=-1)Stacked_value = Value_WeightedStack(attention_scores, value)MUltiHead_output = tf.add(project_final(Stacked_value, MultiHead_output_kerner), MultiHead_output_bias)output = tf.add(output_dense(MUltiHead_output, output_dense_kernel), output_dense_bias)return output

验证

最后我们将分别打印出两个模型前向传播的输出（一个是自己实现的，另一个是TensorFlow实现的），并进行结果比对，看是否相差无几。

test_in = tf.constant([[[0.1, 0.1, 0.1], [0.1, 0.1, 0.1]]], dtype=tf.float32)
test_in0 = tf.constant([[[1, 1, 1], [1, 1, 1]]], dtype=tf.float32)
test_in1 = tf.constant([[[10, 10, 10], [10, 10, 10]]], dtype=tf.float32)
test_in2 = tf.constant([[[100, 100, 100], [100, 100, 100]]], dtype=tf.float32)
test_in3 = tf.constant([[[1000, 1000, 1000], [1000, 1000, 1000]]], dtype=tf.float32)
model = tf.saved_model.load('MultiHeadAttention')
print(model(test_in))
print(model(test_in0))
print(model(test_in1))
print(model(test_in2))
print(model(test_in3))tf.Tensor(
[[[-0.02398136][-0.02398136]]], shape=(1, 2, 1), dtype=float32)
tf.Tensor(
[[[-0.75980777][-0.75980777]]], shape=(1, 2, 1), dtype=float32)
tf.Tensor(
[[[-8.1180725][-8.1180725]]], shape=(1, 2, 1), dtype=float32)
tf.Tensor(
[[[-81.700714][-81.700714]]], shape=(1, 2, 1), dtype=float32)
tf.Tensor(
[[[-817.52716][-817.52716]]], shape=(1, 2, 1), dtype=float32)

test_in = tf.constant([[[0.1, 0.1, 0.1], [0.1, 0.1, 0.1]]], dtype=tf.float32)
test_in0 = tf.constant([[[1, 1, 1], [1, 1, 1]]], dtype=tf.float32)
test_in1 = tf.constant([[[10, 10, 10], [10, 10, 10]]], dtype=tf.float32)
test_in2 = tf.constant([[[100, 100, 100], [100, 100, 100]]], dtype=tf.float32)
test_in3 = tf.constant([[[1000, 1000, 1000], [1000, 1000, 1000]]], dtype=tf.float32)
print(my_model(test_in)())
print(my_model(test_in0)())
print(my_model(test_in1)())
print(my_model(test_in2)())
print(my_model(test_in3)())tf.Tensor(
[[[-0.02398137][-0.02398137]]], shape=(1, 2, 1), dtype=float64)
tf.Tensor(
[[[-0.75980776][-0.75980776]]], shape=(1, 2, 1), dtype=float64)
tf.Tensor(
[[[-8.11807168][-8.11807168]]], shape=(1, 2, 1), dtype=float64)
tf.Tensor(
[[[-81.70071091][-81.70071091]]], shape=(1, 2, 1), dtype=float64)
tf.Tensor(
[[[-817.5271032][-817.5271032]]], shape=(1, 2, 1), dtype=float64)

两个输出结果相差无几，验证成功。

这篇关于用einsum实现MultiHeadAttention前向传播的文章就介绍到这儿，希望我们推荐的文章对编程师们有所帮助！


http://www.chinasem.cn/article/1148314。

相关文章

hdu1043（八数码问题，广搜 + hash（实现状态压缩） ）

hdu1043（八数码问题，广搜 + hash（实现状态压缩） ）

利用康拓展开将一个排列映射成一个自然数，然后就变成了普通的广搜题。 #include<iostream>#include<algorithm>#include<string>#include<stack>#include<queue>#include<map>#include<stdio.h>#include<stdlib.h>#include<ctype.h>#inclu
阅读更多...
【C++】_list常用方法解析及模拟实现

【C++】_list常用方法解析及模拟实现

相信自己的力量，只要对自己始终保持信心，尽自己最大努力去完成任何事，就算事情最终结果是失败了，努力了也不留遗憾。💓💓💓 目录   ✨说在前面 🍋知识点一：什么是list？ •🌰1.list的定义 •🌰2.list的基本特性 •🌰3.常用接口介绍 🍋知识点二：list常用接口 •🌰1.默认成员函数 🔥构造函数（⭐） 🔥析构函数 •🌰2.list对象
阅读更多...
【Prometheus】PromQL向量匹配实现不同标签的向量数据进行运算

【Prometheus】PromQL向量匹配实现不同标签的向量数据进行运算

✨✨ 欢迎大家来到景天科技苑✨✨ 🎈🎈 养成好习惯，先赞后看哦~🎈🎈 🏆 作者简介：景天科技苑 🏆《头衔》：大厂架构师，华为云开发者社区专家博主，阿里云开发者社区专家博主，CSDN全栈领域优质创作者，掘金优秀博主，51CTO博客专家等。 🏆《博客》：Python全栈，前后端开发，小程序开发，人工智能，js逆向，App逆向，网络系统安全，数据分析，Django，fastapi
阅读更多...
让树莓派智能语音助手实现定时提醒功能

让树莓派智能语音助手实现定时提醒功能

最初的时候是想直接在rasa 的chatbot上实现，因为rasa本身是带有remindschedule模块的。不过经过一番折腾后，忽然发现，chatbot上实现的定时，语音助手不一定会有响应。因为，我目前语音助手的代码设置了长时间无应答会结束对话，这样一来，chatbot定时提醒的触发就不会被语音助手获悉。那怎么让语音助手也具有定时提醒功能呢？ 我最后选择的方法是用threading.Time
阅读更多...
Android实现任意版本设置默认的锁屏壁纸和桌面壁纸（两张壁纸可不一致）

Android实现任意版本设置默认的锁屏壁纸和桌面壁纸（两张壁纸可不一致）

客户有些需求需要设置默认壁纸和锁屏壁纸  在默认情况下 这两个壁纸是相同的  如果需要默认的锁屏壁纸和桌面壁纸不一样 需要额外修改 Android13实现 替换默认桌面壁纸： 将图片文件替换frameworks/base/core/res/res/drawable-nodpi/default_wallpaper.*  （注意不能是bmp格式） 替换默认锁屏壁纸： 将图片资源放入vendo
阅读更多...
C#实战|大乐透选号器[6]：实现实时显示已选择的红蓝球数量

C#实战|大乐透选号器[6]：实现实时显示已选择的红蓝球数量

哈喽，你好啊，我是雷工。 关于大乐透选号器在前面已经记录了5篇笔记，这是第6篇； 接下来实现实时显示当前选中红球数量，蓝球数量； 以下为练习笔记。 01 效果演示 当选择和取消选择红球或蓝球时，在对应的位置显示实时已选择的红球、蓝球的数量； 02 标签名称 分别设置Label标签名称为：lblRedCount、lblBlueCount
阅读更多...
Kubernetes PodSecurityPolicy：PSP能实现的5种主要安全策略

Kubernetes PodSecurityPolicy：PSP能实现的5种主要安全策略

Kubernetes PodSecurityPolicy：PSP能实现的5种主要安全策略 1. 特权模式限制2. 宿主机资源隔离3. 用户和组管理4. 权限提升控制5. SELinux配置 💖The Begin💖点点关注，收藏不迷路💖 Kubernetes的PodSecurityPolicy（PSP）是一个关键的安全特性，它在Pod创建之前实施安全策略，确保P
阅读更多...
工厂ERP管理系统实现源码（JAVA）

工厂ERP管理系统实现源码（JAVA）

工厂进销存管理系统是一个集采购管理、仓库管理、生产管理和销售管理于一体的综合解决方案。该系统旨在帮助企业优化流程、提高效率、降低成本，并实时掌握各环节的运营状况。 在采购管理方面，系统能够处理采购订单、供应商管理和采购入库等流程，确保采购过程的透明和高效。仓库管理方面，实现库存的精准管理，包括入库、出库、盘点等操作，确保库存数据的准确性和实时性。 生产管理模块则涵盖了生产计划制定、物料需求计划、
阅读更多...
C++——stack、queue的实现及deque的介绍

C++——stack、queue的实现及deque的介绍

目录 1.stack与queue的实现 1.1stack的实现  1.2 queue的实现 2.重温vector、list、stack、queue的介绍 2.1 STL标准库中stack和queue的底层结构  3.deque的简单介绍 3.1为什么选择deque作为stack和queue的底层默认容器  3.2 STL中对stack与queue的模拟实现 ①stack模拟实现
阅读更多...
基于51单片机的自动转向修复系统的设计与实现

基于51单片机的自动转向修复系统的设计与实现

文章目录 前言资料获取设计介绍功能介绍设计清单具体实现截图参考文献设计获取 前言 💗博主介绍：✌全网粉丝10W+,CSDN特邀作者、博客专家、CSDN新星计划导师，一名热衷于单片机技术探索与分享的博主、专注于 精通51/STM32/MSP430/AVR等单片机设计 主要对象是咱们电子相关专业的大学生，希望您们都共创辉煌！✌💗 👇🏻 精彩专栏 推荐订阅👇🏻 单片机
阅读更多...

Copyright China编程 23002807@qq.com

沪ICP备2023033072号-2