用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实现自动化Word文档样式复制与内容生成

Python实现自动化Word文档样式复制与内容生成

《Python实现自动化Word文档样式复制与内容生成》在办公自动化领域,高效处理Word文档的样式和内容复制是一个常见需求,本文将展示如何利用Python的python-docx库实现... 目录一、为什么需要自动化 Word 文档处理二、核心功能实现：样式与表格的深度复制1. 表格复制（含样式与内容）2
阅读更多...
python获取cmd环境变量值的实现代码

python获取cmd环境变量值的实现代码

《python获取cmd环境变量值的实现代码》：本文主要介绍在Python中获取命令行（cmd）环境变量的值,可以使用标准库中的os模块,需要的朋友可以参考下... 前言全局说明在执行py过程中，总要使用到系统环境变量一、说明1.1 环境:Windows 11 家庭版 24H2 26100.4061
阅读更多...
Python中bisect_left 函数实现高效插入与有序列表管理

Python中bisect_left 函数实现高效插入与有序列表管理

《Python中bisect_left函数实现高效插入与有序列表管理》Python的bisect_left函数通过二分查找高效定位有序列表插入位置,与bisect_right的区别在于处理重复元素时... 目录一、bisect_left 基本介绍1.1 函数定义1.2 核心功能二、bisect_left 与
阅读更多...
VSCode设置python SDK路径的实现步骤

VSCode设置python SDK路径的实现步骤

《VSCode设置pythonSDK路径的实现步骤》本文主要介绍了VSCode设置pythonSDK路径的实现步骤,包括命令面板切换、settings.json配置、环境变量及虚拟环境处理,具有一定... 目录一、通过命令面板快速切换（推荐方法）二、通过 settings.json 配置（项目级/全局）三、
阅读更多...
pandas实现数据concat拼接的示例代码

pandas实现数据concat拼接的示例代码

《pandas实现数据concat拼接的示例代码》pandas.concat用于合并DataFrame或Series,本文主要介绍了pandas实现数据concat拼接的示例代码,具有一定的参考价值,... 目录语法示例：使用pandas.concat合并数据默认的concat：参数axis=0，join=
阅读更多...
java中BigDecimal里面的subtract函数介绍及实现方法

java中BigDecimal里面的subtract函数介绍及实现方法

《java中BigDecimal里面的subtract函数介绍及实现方法》在Java中实现减法操作需要根据数据类型选择不同方法,主要分为数值型减法和字符串减法两种场景,本文给大家介绍java中BigD... 目录Java中BigDecimal里面的subtract函数的意思？一、数值型减法（高精度计算）1.
阅读更多...
C#代码实现解析WTGPS和BD数据

C#代码实现解析WTGPS和BD数据

《C#代码实现解析WTGPS和BD数据》在现代的导航与定位应用中,准确解析GPS和北斗（BD）等卫星定位数据至关重要,本文将使用C#语言实现解析WTGPS和BD数据,需要的可以了解下... 目录一、代码结构概览1. 核心解析方法2. 位置信息解析3. 经纬度转换方法4. 日期和时间戳解析5. 辅助方法二、L
阅读更多...
使用Python和Matplotlib实现可视化字体轮廓(从路径数据到矢量图形)

使用Python和Matplotlib实现可视化字体轮廓(从路径数据到矢量图形)

《使用Python和Matplotlib实现可视化字体轮廓(从路径数据到矢量图形)》字体设计和矢量图形处理是编程中一个有趣且实用的领域,通过Python的matplotlib库,我们可以轻松将字体轮廓... 目录背景知识字体轮廓的表示实现步骤1. 安装依赖库2. 准备数据3. 解析路径指令4. 绘制图形关键
阅读更多...
C/C++中OpenCV 矩阵运算的实现

C/C++中OpenCV 矩阵运算的实现

《C/C++中OpenCV矩阵运算的实现》本文主要介绍了C/C++中OpenCV矩阵运算的实现,包括基本算术运算（标量与矩阵）、矩阵乘法、转置、逆矩阵、行列式、迹、范数等操作,感兴趣的可以了解一下... 目录矩阵的创建与初始化创建矩阵访问矩阵元素基本的算术运算 ➕➖✖️➗矩阵与标量运算矩阵与矩阵运算 (逐元
阅读更多...
C/C++的OpenCV 进行图像梯度提取的几种实现

C/C++的OpenCV 进行图像梯度提取的几种实现

《C/C++的OpenCV进行图像梯度提取的几种实现》本文主要介绍了C/C++的OpenCV进行图像梯度提取的实现,文中通过示例代码介绍的非常详细,对大家的学习或者工作具有一定的参考学习价值,需要的... 目录预www.chinasem.cn备知识1. 图像加载与预处理2. Sobel 算子计算 X 和 Y
阅读更多...

Copyright China编程 23002807@qq.com

沪ICP备2023033072号-2