用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实现批量提取BLF文件时间戳

Python实现批量提取BLF文件时间戳

《Python实现批量提取BLF文件时间戳》BLF（BinaryLoggingFormat）作为Vector公司推出的CAN总线数据记录格式,被广泛用于存储车辆通信数据,本文将使用Python轻松提取... 目录一、为什么需要批量处理 BLF 文件二、核心代码解析：从文件遍历到数据导出1. 环境准备与依赖库
阅读更多...
linux下shell脚本启动jar包实现过程

linux下shell脚本启动jar包实现过程

《linux下shell脚本启动jar包实现过程》确保APP_NAME和LOG_FILE位于目录内,首次启动前需手动创建log文件夹,否则报错,此为个人经验,供参考,欢迎支持脚本之家... 目录linux下shell脚本启动jar包样例1样例2总结linux下shell脚本启动jar包样例1#!/bin
阅读更多...
go动态限制并发数量的实现示例

go动态限制并发数量的实现示例

《go动态限制并发数量的实现示例》本文主要介绍了Go并发控制方法,通过带缓冲通道和第三方库实现并发数量限制,文中通过示例代码介绍的非常详细,对大家的学习或者工作具有一定的参考学习价值,需要的朋友们下面... 目录带有缓冲大小的通道使用第三方库其他控制并发的方法因为go从语言层面支持并发，所以面试百分百会问到
阅读更多...
Go语言并发之通知退出机制的实现

Go语言并发之通知退出机制的实现

《Go语言并发之通知退出机制的实现》本文主要介绍了Go语言并发之通知退出机制的实现,文中通过示例代码介绍的非常详细,对大家的学习或者工作具有一定的参考学习价值,需要的朋友们下面随着小编来一起学习学习吧... 目录1、通知退出机制1.1 进程/main函数退出1.2 通过channel退出1.3 通过cont
阅读更多...
Python实现PDF按页分割的技术指南

Python实现PDF按页分割的技术指南

《Python实现PDF按页分割的技术指南》PDF文件处理是日常工作中的常见需求,特别是当我们需要将大型PDF文档拆分为多个部分时,下面我们就来看看如何使用Python创建一个灵活的PDF分割工具吧... 目录需求分析技术方案工具选择安装依赖完整代码实现使用说明基本用法示例命令输出示例技术亮点实际应用场景扩
阅读更多...
java如何实现高并发场景下三级缓存的数据一致性

java如何实现高并发场景下三级缓存的数据一致性

《java如何实现高并发场景下三级缓存的数据一致性》这篇文章主要为大家详细介绍了java如何实现高并发场景下三级缓存的数据一致性,文中的示例代码讲解详细,感兴趣的小伙伴可以跟随小编一起学习一下... 下面代码是一个使用Java和Redisson实现的三级缓存服务，主要功能包括：1.缓存结构：本地缓存：使
阅读更多...
如何在Java Spring实现异步执行(详细篇)

如何在Java Spring实现异步执行(详细篇)

《如何在JavaSpring实现异步执行(详细篇)》Spring框架通过@Async、Executor等实现异步执行,提升系统性能与响应速度,支持自定义线程池管理并发,本文给大家介绍如何在Sprin... 目录前言1. 使用 @Async 实现异步执行1.1 启用异步执行支持1.2 创建异步方法1.3 调用
阅读更多...
Spring Boot配置和使用两个数据源的实现步骤

Spring Boot配置和使用两个数据源的实现步骤

《SpringBoot配置和使用两个数据源的实现步骤》本文详解SpringBoot配置双数据源方法,包含配置文件设置、Bean创建、事务管理器配置及@Qualifier注解使用,强调主数据源标记、代... 目录Spring Boot配置和使用两个数据源技术背景实现步骤1. 配置数据源信息2. 创建数据源Be
阅读更多...
在MySQL中实现冷热数据分离的方法及使用场景底层原理解析

在MySQL中实现冷热数据分离的方法及使用场景底层原理解析

《在MySQL中实现冷热数据分离的方法及使用场景底层原理解析》MySQL冷热数据分离通过分表/分区策略、数据归档和索引优化,将频繁访问的热数据与冷数据分开存储,提升查询效率并降低存储成本,适用于高并发... 目录实现冷热数据分离1. 分表策略2. 使用分区表3. 数据归档与迁移在mysql中实现冷热数据分
阅读更多...
linux批量替换文件内容的实现方式

linux批量替换文件内容的实现方式

《linux批量替换文件内容的实现方式》本文总结了Linux中批量替换文件内容的几种方法,包括使用sed替换文件夹内所有文件、单个文件内容及逐行字符串,强调使用反引号和绝对路径,并分享个人经验供参考... 目录一、linux批量替换文件内容 二、替换文件内所有匹配的字符串 三、替换每一行中全部str1为st
阅读更多...

Copyright China编程 23002807@qq.com

沪ICP备2023033072号-2