用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。

相关文章

MySQL中查找重复值的实现

MySQL中查找重复值的实现

《MySQL中查找重复值的实现》查找重复值是一项常见需求,比如在数据清理、数据分析、数据质量检查等场景下,我们常常需要找出表中某列或多列的重复值,具有一定的参考价值,感兴趣的可以了解一下... 目录技术背景实现步骤方法一：使用GROUP BY和HAVING子句方法二：仅返回重复值方法三：返回完整记录方法四：
阅读更多...
IDEA中新建/切换Git分支的实现步骤

IDEA中新建/切换Git分支的实现步骤

《IDEA中新建/切换Git分支的实现步骤》本文主要介绍了IDEA中新建/切换Git分支的实现步骤,通过菜单创建新分支并选择是否切换,创建后在Git详情或右键Checkout中切换分支,感兴趣的可以了... 前提：项目已被Git托管1、点击上方栏Git->NewBrancjsh...2、输入新的分支的
阅读更多...
Spring事务传播机制最佳实践

Spring事务传播机制最佳实践

《Spring事务传播机制最佳实践》Spring的事务传播机制为我们提供了优雅的解决方案,本文将带您深入理解这一机制,掌握不同场景下的最佳实践,感兴趣的朋友一起看看吧... 目录1. 什么是事务传播行为2. Spring支持的七种事务传播行为2.1 REQUIRED（默认）2.2 SUPPORTS2
阅读更多...
Python实现对阿里云OSS对象存储的操作详解

Python实现对阿里云OSS对象存储的操作详解

《Python实现对阿里云OSS对象存储的操作详解》这篇文章主要为大家详细介绍了Python实现对阿里云OSS对象存储的操作相关知识,包括连接,上传,下载,列举等功能,感兴趣的小伙伴可以了解下... 目录一、直接使用代码二、详细使用1. 环境准备2. 初始化配置3. bucket配置创建4. 文件上传到os
阅读更多...
关于集合与数组转换实现方法

关于集合与数组转换实现方法

《关于集合与数组转换实现方法》：本文主要介绍关于集合与数组转换实现方法,具有很好的参考价值,希望对大家有所帮助,如有错误或未考虑完全的地方,望不吝赐教... 目录1、Arrays.asList()1.1、方法作用1.2、内部实现1.3、修改元素的影响1.4、注意事项2、list.toArray()2.1、方
阅读更多...
使用Python实现可恢复式多线程下载器

使用Python实现可恢复式多线程下载器

《使用Python实现可恢复式多线程下载器》在数字时代,大文件下载已成为日常操作,本文将手把手教你用Python打造专业级下载器,实现断点续传,多线程加速,速度限制等功能,感兴趣的小伙伴可以了解下... 目录一、智能续传：从崩溃边缘抢救进度二、多线程加速：榨干网络带宽三、速度控制：做网络的好邻居四、终端交互
阅读更多...
java实现docker镜像上传到harbor仓库的方式

java实现docker镜像上传到harbor仓库的方式

《java实现docker镜像上传到harbor仓库的方式》：本文主要介绍java实现docker镜像上传到harbor仓库的方式,具有很好的参考价值,希望对大家有所帮助,如有错误或未考虑完全的地... 目录1. 前 言2. 编写工具类2.1 引入依赖包2.2 使用当前服务器的docker环境推送镜像2.2
阅读更多...
C++20管道运算符的实现示例

C++20管道运算符的实现示例

《C++20管道运算符的实现示例》本文简要介绍C++20管道运算符的使用与实现,文中通过示例代码介绍的非常详细,对大家的学习或者工作具有一定的参考学习价值,需要的朋友们下面随着小编来一起学习学习吧... 目录标准库的管道运算符使用自己实现类似的管道运算符我们不打算介绍太多，因为它实际属于c++20最为重要的
阅读更多...
Java easyExcel实现导入多sheet的Excel

Java easyExcel实现导入多sheet的Excel

《JavaeasyExcel实现导入多sheet的Excel》这篇文章主要为大家详细介绍了如何使用JavaeasyExcel实现导入多sheet的Excel,文中的示例代码讲解详细,感兴趣的小伙伴可... 目录1.官网2.Excel样式3.代码1.官网easyExcel官网2.Excel样式3.代码
阅读更多...
python实现对数据公钥加密与私钥解密

python实现对数据公钥加密与私钥解密

《python实现对数据公钥加密与私钥解密》这篇文章主要为大家详细介绍了如何使用python实现对数据公钥加密与私钥解密,文中的示例代码讲解详细,感兴趣的小伙伴可以跟随小编一起学习一下... 目录公钥私钥的生成使用公钥加密使用私钥解密公钥私钥的生成这一部分，使用python生成公钥与私钥，然后保存在两个文
阅读更多...

Copyright China编程 23002807@qq.com

沪ICP备2023033072号-2