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

相关文章

Python如何实现PDF隐私信息检测

Python如何实现PDF隐私信息检测

《Python如何实现PDF隐私信息检测》随着越来越多的个人信息以电子形式存储和传输,确保这些信息的安全至关重要,本文将介绍如何使用Python检测PDF文件中的隐私信息,需要的可以参考下... 目录项目背景技术栈代码解析功能说明运行结php果在当今，数据隐私保护变得尤为重要。随着越来越多的个人信息以电子形
阅读更多...
使用 sql-research-assistant进行 SQL 数据库研究的实战指南(代码实现演示)

使用 sql-research-assistant进行 SQL 数据库研究的实战指南(代码实现演示)

《使用sql-research-assistant进行SQL数据库研究的实战指南(代码实现演示)》本文介绍了sql-research-assistant工具,该工具基于LangChain框架,集... 目录技术背景介绍核心原理解析代码实现演示安装和配置项目集成LangSmith 配置（可选）启动服务应用场景
阅读更多...
使用Python快速实现链接转word文档

使用Python快速实现链接转word文档

《使用Python快速实现链接转word文档》这篇文章主要为大家详细介绍了如何使用Python快速实现链接转word文档功能,文中的示例代码讲解详细,感兴趣的小伙伴可以跟随小编一起学习一下... 演示代码展示from newspaper import Articlefrom docx import
阅读更多...
前端原生js实现拖拽排课效果实例

前端原生js实现拖拽排课效果实例

《前端原生js实现拖拽排课效果实例》：本文主要介绍如何实现一个简单的课程表拖拽功能,通过HTML、CSS和JavaScript的配合,我们实现了课程项的拖拽、放置和显示功能,文中通过实例代码介绍的... 目录1. 效果展示2. 效果分析2.1 关键点2.2 实现方法3. 代码实现3.1 html部分3.2
阅读更多...
Java深度学习库DJL实现Python的NumPy方式

Java深度学习库DJL实现Python的NumPy方式

《Java深度学习库DJL实现Python的NumPy方式》本文介绍了DJL库的背景和基本功能,包括NDArray的创建、数学运算、数据获取和设置等,同时,还展示了如何使用NDArray进行数据预处理... 目录1 NDArray 的背景介绍1.1 架构2 JavaDJL使用2.1 安装DJL2.2 基本操
阅读更多...
最长公共子序列问题的深度分析与Java实现方式

最长公共子序列问题的深度分析与Java实现方式

《最长公共子序列问题的深度分析与Java实现方式》本文详细介绍了最长公共子序列（LCS）问题,包括其概念、暴力解法、动态规划解法,并提供了Java代码实现,暴力解法虽然简单,但在大数据处理中效率较低,... 目录最长公共子序列问题概述问题理解与示例分析暴力解法思路与示例代码动态规划解法DP 表的构建与意义动
阅读更多...
java父子线程之间实现共享传递数据

java父子线程之间实现共享传递数据

《java父子线程之间实现共享传递数据》本文介绍了Java中父子线程间共享传递数据的几种方法,包括ThreadLocal变量、并发集合和内存队列或消息队列,并提醒注意并发安全问题... 目录通过 ThreadLocal 变量共享数据通过并发集合共享数据通过内存队列或消息队列共享数据注意并发安全问题总结在 J
阅读更多...
SpringBoot+MyBatis-Flex配置ProxySQL的实现步骤

SpringBoot+MyBatis-Flex配置ProxySQL的实现步骤

《SpringBoot+MyBatis-Flex配置ProxySQL的实现步骤》本文主要介绍了SpringBoot+MyBatis-Flex配置ProxySQL的实现步骤,文中通过示例代码介绍的非常详... 目录 目标 步骤 1：确保 ProxySQL 和 mysql 主从同步已正确配置ProxySQL 的
阅读更多...
JS 实现复制到剪贴板的几种方式小结

JS 实现复制到剪贴板的几种方式小结

《JS实现复制到剪贴板的几种方式小结》本文主要介绍了JS实现复制到剪贴板的几种方式小结,包括ClipboardAPI和document.execCommand这两种方法,具有一定的参考价值,感兴趣的... 目录一、Clipboard API相关属性方法二、document.execCommand优点：缺点：
阅读更多...
nginx部署https网站的实现步骤(亲测)

nginx部署https网站的实现步骤(亲测)

《nginx部署https网站的实现步骤(亲测)》本文详细介绍了使用Nginx在保持与http服务兼容的情况下部署HTTPS,文中通过示例代码介绍的非常详细,对大家的学习或者工作具有一定的参考学习价值... 目录步骤 1：安装 Nginx步骤 2：获取 SSL 证书步骤 3：手动配置 Nginx步骤 4：测
阅读更多...

Copyright China编程 23002807@qq.com

沪ICP备2023033072号-2