2021SC@SDUSC
lingvo.core.favor_attention包
查看导入的compat中基于tensorflow的方法
from lingvo import compat as tf
from lingvo.core import py_utils
compat: 兼容性. 用于进行tensorflow的兼容处理
内置方法:
- 数据集处理
class _CacheDataset(dataset_ops.UnaryUnchangedStructureDataset):
- 基于无状态缓存实现缓存数据集中的元素
def stateless_cache_dataset(filename=""):
- 基于无状态打乱实现随机洗牌数据集的元素
def stateless_shuffle_dataset(buffer_size, seed=None, reshuffle_each_iteration=None)
- 匹配一个或多个glob模式的所有文件的数据集
def stateless_list_files(file_pattern, shuffle=None, seed=None)
- 基于无状态缓存实现缓存数据集中的元素
- 使用附加的错误消息重写 tf. compatible .v1.variable_scope
class variable_scope(tf1.variable_scope)
包中方法
causal_numerator(qs, ks, vs)
函数作用: 计算非归一化的条件注意力向量FAVOR.
参数:
- qs: 形状为 [L, B, H, M] 的 query_prime 张量
- ks: 形状为 [L,B,H,M] 的 key_prime 张量
- vs: 形状为 [L,B,H,D] 的数值型张量
返回值:
计算非归一化的条件注意力向量FAVOR
源码
def causal_numerator(qs, ks, vs):
result = []
sums = tf.zeros_like(tf.einsum("ijk,ijl->ijkl", ks[0], vs[0]))
for index in range(qs.shape[0]):
sums = sums + tf.einsum("ijk,ijl->ijkl", ks[index], vs[index])
result.append(tf.einsum("ijkl,ijk->ijl", sums, qs[index])[None, ...])
result = tf.concat(result, axis=0)
def grad(res_grad):
grads = tf.zeros_like(tf.einsum("ijk,ijl->ijkl", ks[0], vs[0]))
gr_sums = sums
q_grads = []
k_grads = []
v_grads = []
通过计算梯度得到张量
for index in range(qs.shape[0] - 1, -1, -1):
q_grads.append(
tf.einsum("ijkl,ijl->ijk", gr_sums, res_grad[index])[None, ...])
grads = grads + tf.einsum("ijk,ijl->ijkl", qs[index], res_grad[index])
k_grads.append(tf.einsum("ijkl,ijl->ijk", grads, vs[index])[None, ...])
v_grads.append(tf.einsum("ijkl,ijk->ijl", grads, ks[index])[None, ...])
gr_sums = gr_sums - tf.einsum("ijk,ijl->ijkl", ks[index], vs[index])
q_grads = tf.concat(q_grads[::-1], axis=0)
k_grads = tf.concat(k_grads[::-1], axis=0)
v_grads = tf.concat(v_grads[::-1], axis=0)
return q_grads, k_grads, v_grads
return result, grad
在代码中发现有趣的事情----方法中定义了一个新的方法
在python中这种写法称为闭包,可以使用外层函数接收的参数,Python 里面实现装饰器的基础,很多语言都支持这个写法。有时为了保证一个方法所实现功能的完整性,避免代码过于零散也可以这么写,实现多线程多进程的时候也可以直接定义一个内部函数传给 target,看起来会好看点
关于Favor_attention
传统注意力机制的分解可以得到线性(而非二次)空间复杂度隐式注意力矩阵。同样,通过分解可以获得线性时间的注意力机制。原有的方式是注意力矩阵与value输入相乘得到最终结果,但在分解注意力矩阵之后,可以重新排列矩阵乘法来近似常规注意力机制的结果,而无需显式地构建二次方尺寸的注意力矩阵。最终生成了新算法 FAVOR+。
上述分析是双向注意力(即非因果注意力),并没有区分 past 和 future的部分。那么如何做到让输入序列中只注意到其中一部分,即单向的因果注意力?只需使用前缀和计算(prefix-sum computation),在计算过程只存储矩阵计算的运行总数,而不存储完整的下三角常规注意力矩阵。
参考资料–知乎–Performer带头反思Attention
layers_with_attention
随机残层指的是随机掉落残余分支, 该观点最初在 “Deep Networks with Stochastic Depth” 中的CNN被提出.
class TransformerAttentionLayer(base_layer.BaseLayer)
模型结构: Multi-headed attention → Add&Norm
该类实现了转换器层的第一个子层。输入是首先使用多头(自)注意力处理。注意力层的输出与残余连接相结合。最后, 输出使用层归一化进行规范化。
该层可以在五种情况下使用:
- 多头自注意,即attention键 (源载体),attention值 (上下文向量) 和查询来自相同的前一个输出层 “query_vec”. 这是transformer层编码器的一般用例.
- mask多头自注意,其中attention键,attention值和查询都来自相同的前一层输出,但右激活被屏蔽,以阻止未来的信息流. 这是解码器自注意transformer层的用例。用is_mask标志该层是否可以被激活.
- 多头注意,其中attention键和attention值 ’ source_vecs ‘,来自不同的资源 (编码器的输出) 查询’query_vec’,来自前一层输出 (解码器)。这就对应了标准的注意机制,解码器处理编码器输出。
- Multi-Headed Attention,注意值“context_vecs”来源于不同的从不同的查询和键值,例如条件注意力机制,键和查询是条件编码,值是条件编码译码器的状态。
- masked多头自注意机制,其中attention键,attention值和查询都来自相同的前一层输出,但当前位置的激活被屏蔽,以减少高的自相似性。这是非自回归解码器的用例解码器自注意transformer层。可以通过设置is_mask来激活----设置mask_type=“eye”。
关于mask
mask可以理解成遮罩、面具,作用是帮助我们“遮挡”掉我们不需要的东西,即让被遮挡的东西不影响我们的attention过程。
在forward的时候,有两个mask参数可以设置:
key_padding_mask
每一个batch的每一个句子的长度一般是不可能完全相同的,所以我们会使用padding把一些空缺补上。而这里的这个key_padding_mask是用来“遮挡”这些padding的。
这个mask是二元(binary)的,也就是说,它是一个矩阵和我们key的大小是一样的,里面的值是1或0,我们先取得key中有padding的位置,然后把mask里相应位置的数字设置为1,这样attention就会把key相应的部分变为"-inf".
attn_mask
这个mask经常是用来遮挡“正确答案”的:
假如你想要用这个模型每次预测下一个单词,我们每一个位置的attention输出是怎么得来的?是不是要看一遍整个序列,然后每一个单词都计算一个attention weight?那也就是说,你在预测第5个词的时候,你其实会看到整个序列,这样的话你在预测之前不就已经知道第5个单词是什么了,这就是作弊了。
我们不想让模型作弊,因为在真实使用这个模型去预测的时候,我们是没有整个序列的信息的。那么怎么办?那就让第5个单词的attention weight=0吧,即声明:我不想看这个单词,我的注意力一点也别分给它。
参考「狗狗狗大王」原文链接:https://blog.csdn.net/weixin_41811314/article/details/106804906
方法
def FProp(self, theta, query_vec, source_paddings, source_vecs=None, query_segment_id=None, source_segment_id=None, context_vecs=None, **kwargs)
是 transformer 的注意力机制 残差 标准化层
参数列表:
- theta:
.NestedMap
对象,其中包含此层和的权重值和其子层。 - query_vec: [target_time, target_batch, dim]
- source_paddings: [source_time, source_batch]
- source_vecs: [source_time, source_batch, dim].
- query_segment_id: [target_time, target_batch]
- source_segment_id: [source_time, source_batch]
- context_vecs: [source_time, target_batch, dim]
- **kwargs: 可以是注意层的可选参数,如注意投射指数张量。
源码
def FProp(self,
theta,
query_vec,
source_paddings,
source_vecs=None,
query_segment_id=None,
source_segment_id=None,
context_vecs=None,
**kwargs):
p = self.params
unnormalized_query_vec = query_vec
if p.pre_layer_norm:
query_vec = self.layer_norm.FProp(theta.layer_norm, query_vec)
对于自注意力机制: keys = queries.
if source_vecs is None:
source_vecs = query_vec
source_segment_id = query_segment_id
Inter/self-attention: keys = values/contexts.
if context_vecs is None:
context_vecs = source_vecs
target_time, target_bs, query_dim = py_utils.GetShape(query_vec, 3)
if p.is_masked:
assert source_vecs is not None
query_vec = py_utils.with_dependencies([
py_utils.assert_shape_match(
tf.shape(source_vecs), tf.shape(query_vec))
], query_vec)
- 准备 self-attention 的 mask
- 填充被补全, 掩盖的时间索引
- 接收padding 权重为1.0。
if p.mask_type == 'future':
padding = py_utils.CausalSelfAttenPadding(
target_time, dtype=py_utils.FPropDtype(p))
elif p.mask_type == 'eye':
padding = tf.eye(target_time, target_time, dtype=py_utils.FPropDtype(p))
elif p.mask_type == 'ngram': # Maybe apply N-gram mask.
assert p.mask_ngram_order
padding = 1.0 - tf.linalg.band_part(
tf.ones([target_time, target_time], dtype=py_utils.FPropDtype(p)),
tf.minimum(p.mask_ngram_order - 1, target_time - 1), 0)
# [time, batch, time]
causal_padding = tf.tile(tf.expand_dims(padding, 1), [1, target_bs, 1])
causal_padding = tf.reshape(causal_padding, [-1, target_time])
else:
causal_padding = None
键值对
packed_src = self.atten.PackSource(
theta=theta.atten,
source_vecs=source_vecs, # keys
source_contexts=context_vecs, # values
source_padding=source_paddings,
source_segment_id=source_segment_id)
if query_segment_id is not None:
query_segment_id = tf.reshape(query_segment_id, [-1])
ctx_vec, atten_prob, _ = self.atten.ComputeContextVectorWithSource(
theta=theta.atten,
packed_src=packed_src,
query_vec=tf.reshape(query_vec, [-1, query_dim]),
per_step_source_padding=causal_padding,
query_segment_id=query_segment_id,
**kwargs)
ctx_vec = self.residual_dropout.FProp(theta.residual_dropout, ctx_vec)
input_to_add = (
unnormalized_query_vec if p.add_unnormalized_input else query_vec)
input_after_sublayer = tf.reshape(
ctx_vec,
[
target_time,
target_bs,
-1 # Either projected or not.
])
if p.residual_function is None:
h = input_to_add + input_after_sublayer
else:
h = self.residual_function.FProp(theta.residual_function, input_to_add,
input_after_sublayer)
if not p.pre_layer_norm:
h = self.layer_norm.FProp(theta.layer_norm, h)
atten_prob = tf.reshape(
atten_prob,
[target_time, target_bs,
self._GetSourceLength(source_paddings)])
return h, atten_prob