最后更新于
最后更新于
本文介绍XLNet的代码的第二部分,需要首先阅读,读者阅读前需要了解XLNet的原理,不熟悉的读者请先阅读。
目录
书接上文,前面我们提到:如果忽略多设备(GPU)训练的细节,train的代码结构其实并不复杂,它大致可以分为3部分:
调用data_utils.get_input_fn得到train_input_fn
调用single_core_graph构造XLNet网络
使用session运行fetches进行训练
首先我们来看一下train函数的主要代码,为了简单,我们这里只分析单个GPU的代码路径而忽略多GPU的情况:
我们在计算当前句子时会用到之前的96个Token作为上下文,因此需要创建mems:
在我的例子里它是创建了长度为6的list(因为我测试用的XLNet只有6层),每一个元素的shape是[96, 8, 1024],96是mem的长度,8是batch_size,1024是隐单元个数。
上面的train函数的训练过程非常的简单,接下面我们详细的来介绍train函数处理Dataset输入以及构建计算图的代码。
返回值知道了,我们再来看这个函数的参数:
tfrecord_dir 训练数据目录,如果有多个用逗号”,”分开,这里是’traindata/tfrecords’,参考第一部分。
split 这里是”train”表示训练数据。
bsz_per_host 每个host(机器)的batch大小,这里是8。
seq_len 句子长度,这里是128。
reuse_len 句子中前面不变的部分,这里是64,请参考第一部分。
bi_data 是否双向的模型,这里是True。
num_hosts=1 有多少台机器,这里是1。
num_core_per_host=1 每天机器有多少个设备(GPU),这里是1。
perm_size=None, 排列打散的长度,这里是32,后面会讲到。
mask_alpha=None, 这里是6
mask_beta=None, 这里是1
uncased=False, 是否不区分大小,这里是True,表示不区分大小写
num_passes=None, 训练的趟数,这里是1表示这是第一趟训练
use_bfloat16=False, 是否使用16位表示的浮点数,这里是False
num_predict=None: 用于预测的Token的个数,这里是21
下面我们来看这个函数的代码,请仔细阅读其中的注释。
这个函数主要的代码就是遍历目录下的json文件,然后找到对应的tfrecord文件,放到record_info里,同时返回input_fn函数,如果执行input_fn函数,则它会调用get_dataset函数返回Dataset,下面我们来看get_dataset函数。
代码如下:
主要的代码都在parser函数里,它定义了怎么读取我们在第一部分代码里生成的tfrecord文件,然后进行排列打散(我们还没有具体介绍的_local_perm函数),最后把处理的结果放到example里。接着使用parse_files_to_dataset来得到Dataset,我们来看这个函数:
这个函数比较难懂,这里介绍一个小技巧。如果读者对比过PyTorch和Tensorflow就会感觉使用PyTorch调试会简单很多,原因就是PyTorch是动态的计算图,因此我们把两个Tensor一相加,结果马上就出来了;但是Tensoflow是静态图,我们定义了两个Tensor相加后得到的不是它们的计算结果,而是一个Operation,我们还要用session来run它才能得到结果。如果计算简单还好,一个复杂的的函数经过一系列变换后就完全不知道它的shape是什么了。因此调试PyTorch的代码就像调试普通的Python代码一样;而调试Tensorflow的代码就像”阅读”Pyton代码一样——你看不到执行的结果。不过还有Tensorflow引入了Eager Execution。但是很多代码(包括这里的XLNet)还是习惯用之前的静态构造方法。不过没有关系,如果我们遇到一个函数看不到,那么我们可以对这个函数进行Eager Execution来调试它。比如我们如果看不到_local_perm,则我们可以这样来调试它:
其中seq_len表示输入的长度,perm_size表示排列打散的长度。inputs表示输入,targets表示输出,其中4和3是特殊的SEP和CLS。is_masked表示某个输入是否是MASK(用于预测的),上面的例子里第5和6个位置、第13和14个位置都是Masked,也就是需要模型预测计算loss的位置。
接下来我们就可以增加断点单步执行从而来了解这个函数的作用了。
根据上面的输入,首先用tf.range生成[0, 15]的序列,然后第二行代码首先把它reshape成[2, 8],然后transpose成[8, 2],从而得到:
然后使用tf.random_shuffle对第一列进行随机打散,得到:
最后transpose然后reshape成长度为16的向量:
总结一下代码的效果:把长度为seq_len(=16)的向量分成seq_len/perm_size(=2)段,每段进行随机打散。
1.首先是得到non_func_tokens,所谓的non_func_tokens是指SEP和CLS之外的”正常”的Token。
根据前面的输入inputs,第7个位置为SEP,15和16为SEP和CLS,因此这三个位置为False,其余都为True。
2.然后是non_mask_tokens,non_mask_tokens指的是”正常”的并且没有被Mask的Token。
因此如果一个Token是non_mask_tokens,那么它首先是”正常”的Token(non_func_tokens为True),然后还得没有被Mask(is_masked为False)。
3.masked_or_func_tokens,它和non_mask_tokens相反,包括Masked的Token和SEP与CLS。
tf.where函数的作用等价于:
这样得到的rev_index为:如果某个位置是非Mask的,则其值为-1;反正如果某个位置是Mask(5,6,13,14)或者为特殊的SEP/CLS(7,15,16),则值为前面随机生成的下标。
我们看到,is_masked为True的下标对应的target_mask为1,其余为0。
因为target_tokens不能看到自己,因此它的值就是rev_index,而其余的加1(如果原来是-1,那么变成0,否则就是排列下标加一)。
上面的代码”self_rev_index[:, None] <= rev_index[None, :]”我们来详细分析一下。
首先self_rev_index是一个长度为16的Tensor,我们使用self_rev_index[:, None]把它变成(16,1)的Tensor。
rev_index也是长度为的Tensor,rev_index[None, :]变成了(1,16)的Tensor。然后它们比较(“<=”)时会使用Broadcasting,我们可以了解为self_rev_index[:, None]变成了(16,16)的,每行都是相同的:
类似的,rev_index[None, :]在Broadcasting之前为:
而Broadcasting之后变成:
因此,最终得到的perm_mask(i,j)=1,则表示i不能attend to j。有两种情况i能attend to j:i的排列下标大于j(后面的可以attend to前面的);j没有被Mask在i也可以attend to j。而i不能attend to j需要同时满足:i<=j并且j被Mask了我们来看一个例子:perm_mask(3,4)=1,因为第3个Token的排列下标是2,第4个的排列下标是3,所以满足”2<3”。请读者检查确认一下j确实是被Mask了。
对于常规的语言模型来说,我们是预测下一个词,而XLNet是根据之前的状态和当前的位置预测被MASK的当前词。所以真正的new_targets要前移一个。
最终的返回值:
如果读者还没有完全明白,我们再来把代码和论文的描述对比一下。论文在Partial Prediction部分提到:为了提高效率,我们只预测排列最后几个词。比如假设总共4个词,并且假设某次随机排列的顺序为3→2→4→1$ 3\rightarrow 2\rightarrow 4\rightarrow 1 $,我们假设预测(Mask)最后两个词,也就是预测第4和第1个词。那么1可以attend to [2,3,4];4可以attend to [2,3]。而非Mask(预测)的第2和3个词使用普通的Self-Attention,也就是可以互相看到所有的非Mask词(包括自己),因此2可以attend to [2,3],3也可以attend to [2,3]。
上面按照论文我们随机排列,然后选择最后的几个词作为Mask;而前面的代码我们先已经选定了要Mask的值,然后再排列,这就可能出现非Mask值的排列下标有可能Mask的值,比如假设我们选定Mask第2个和第3个词,但是随机的排列为:
也就是说第2个词在排列里的顺序是1,第3个词是2。那么按照论文应该是预测第1个和第4个词。那怎么办呢?代码使用了一个tricky,把所有的非Mask的词(第1个和第4个都变成-1),而Mask的不变(总是大于0的),因此Mask的词就排在后面了。非Mask的词互相都可以attend to但是非Mask的词不能attend to Mask的词。Mask的词可以attend to 非Mask的词而且后面的Mask的词也能attend to 前面的Mask的词。比如上面的例子2和3都是Mask的词,因为3在2后面,所以3可以attend to 2,但2不能attend to 3。同时,Mask的词不能attend to 自己(否则就是用自己预测自己了)。
如果读者还是没有理解这个函数也没有太多关系,但是至少要知道这个函数的干什么的,也就是输入输出是什么,具体怎么实现的可以当成黑盒。
总结一下,_local_perm返回的值:
perm_mask,64x64,表示经过重新排列后第i个token能否attend to 第j个token,1表示不能attend
target,64,表示真实的目标值,之前生成的target是预测下一个词,但是XLNet是预测当前词
target_mask,64,哪些地方是Mask的(需要预测的)
input_k, 64,content stream的初始值
input_q, 64, 哪些位置是需要计算loss的,如果不计算loss,也就不计算Query Stream。
这个函数的作用构建XLNet计算图的,是我们需要重点理解的部分。不过这个函数没几行代码:
它调用get_model_fn得到model function,然后调用这个函数返回total_loss, new_mems, grads_and_vars这3个Operation,下面我们来看get_model_fn函数。
这个函数主要是调用function_builder.get_loss得到total_loss, new_mems,然后计算梯度并且返回。
XLNetConfig包含某个checkpoint特定的超参数,也就是pretraining和finetuing都相同的超参数。它包括:
n_layer: int, XLNet的层数,这里是6。
d_model: int, 隐单元个数,这里是1024。
n_head: int, attention head的个数,这里是16。
d_head: int, 每个attention head的大小,要求n_head*d_head=d_model,这里是64。
d_inner: int, 全连接网络的隐单元个数,这里是4096。
ff_activation: str, “relu”或者”gelu”。
untie_r: bool, 是否不同层不共享bias,这里为True,也就是每层都有独立的bias。
n_token: int, 词典大小,这里是32000。
这是Pretraining的一些超参数:
这是真正定义XLNet模型的类。它定义XLNet的代码都在构造函数里,其余的一些函数都是一些getter类函数,比如get_sequence_output可以拿到最后一层的隐状态。我们首先来看它的构造函数的参数。
xlnet_config: XLNetConfig,XLNet模型结构的超参数,比如层数,head数量等等
run_config: RunConfig,运行时的超参数,包括dropout、初始范围等等。
input_ids: int32 Tensor,shape是[len, bsz], 输入token的ID
seg_ids: int32 Tensor,shape是[len, bsz], 输入的segment ID
input_mask: float32 Tensor,shape是[len, bsz], 输入的mask,0是真正的tokens而1是padding的
mems: list,每个元素是float32 Tensors,shape是[mem_len, bsz, d_model], 上一个batch的memory
perm_mask: float32 Tensor,shape是[len, len, bsz]。
如果perm_mask[i, j, k] = 0,则batch k的第i个Token可以attend to j
如果perm_mask[i, j, k] = 1, 则batch k的第i个Token不可以attend to j
如果是None,则每个位置都可以attend to 所有其它位置(包括自己)
target_mapping: float32 Tensor,shape是[num_predict, len, bsz]
如果target_mapping[i, j, k] = 1,则batch k的第i个要预测的是第j个Token,这是一种one-hot表示
只是在pretraining的partial prediction时使用,finetuning时设置为None
inp_q: float32 Tensor,shape是[len, bsz]
需要计算loss的(Mask的位置)为1,不需要的值为0,只在pretraining使用,finetuning时应为None
注意:这里的input_mask不是我们之前介绍的is_masked,is_masked是表示这个位置的Token是被预测的。而这里的input_mask其实指的是padding,我们这里是Pretraining,所有的Token都是真实的,因此传入的为None,后面的代码会把None当成没有padding处理。
上面的构造函数核心的代码其实只有一行:
它的构造函数的参数是前面XLNetModel的构造函数传过来的,不过我们还是列举一下。
inp_k: int32 Tensor,shape是[len, bsz], 输入token的ID
seg_ids: int32 Tensor,shape是[len, bsz], 输入的segment ID
input_mask: float32 Tensor,shape是[len, bsz], 输入的mask,0是真正的tokens而1是padding的
mems: list,每个元素是float32 Tensors,shape是[mem_len, bsz, d_model], 上一个batch的memory
perm_mask: float32 Tensor,shape是[len, len, bsz]。
如果perm_mask[i, j, k] = 0,则batch k的第i个Token可以attend to j
如果perm_mask[i, j, k] = 1, 则batch k的第i个Token不可以attend to j
如果是None,则每个位置都可以attend to 所有其它位置(包括自己)
target_mapping: float32 Tensor,shape是[num_predict, len, bsz]
如果target_mapping[i, j, k] = 1,则batch k的第i个要预测的是第j个Token,这是一种one-hot表示
只是在pretraining的partial prediction时使用,finetuning时设置为None
inp_q: float32 Tensor,shape是[len, bsz]
需要计算loss的(Mask的位置)为1,不需要的值为0,只在pretraining使用,finetuning时应为None
n_layer: int, XLNet的层数,这里是6。
d_model: int, 隐单元个数,这里是1024。
n_head: int, attention head的个数,这里是16。
d_head: int, 每个attention head的大小,要求n_head*d_head=d_model,这里是64。
d_inner: int, 全连接网络的隐单元个数,这里是4096。
ff_activation: str, “relu”或者”gelu”。
untie_r: bool, 是否不同层不共享bias,这里为True,也就是每层都有独立的bias。
n_token: int, 词典大小,这里是32000。
is_training: bool, 是否是Training
use_tpu: bool, 是否使用TPU
use_bfloat16: bool, 是否用bfloat16替代float32
dropout: float, dropout大小.
dropatt: float, attention概率的dropout
init: str, 初始化方法,值为”normal”或者”uniform”
init_range: float, 均匀分布的范围,[-init_range, init_range]。只有init=”uniform”时有效
init_std: float, 正态分布的方程。只有init=”normal”时有效
mem_len: int, cache的token个数
reuse_len: int, 当前batch里reuse的数量,参考第一部分。
bi_data: bool, 是否双向处理输入,通常在pretraining设置为True,而finetuning为False
clamp_len: int, clamp掉相对距离大于clamp_len的attention,-1表示不clamp
same_length: bool, 是否对于每个Token使用相同的attention长度
summary_type: str, “last”, “first”, “mean”和”attn”。怎么把最后一层的多个向量整合成一个向量
initializer: 初始化器
scope: 计算图的scope
下面我们来分段阅读这个最重要的函数。
第一段
上面的代码注意是定义r_w_bias和r_r_bias,以及读取一些超参数。
第二段
下面来看一下non_tgt_mask,为了简单,我们假设qlen是4, mlen是3,前两行的结果为:
attn_mask是(qlen, qlen+mlen, batch, 1),它和non_tgt_mask[:,:,None,None]相加。它的作用是让Token能attend to 自己,除后面的对角线外,non_tgt_mask等于attn_mask,而对角线的位置由1变成了0,从而让它可以attend自己。注意:在XLNet的代码里,0表示可以attend to,而1表示不能attend to。non_tgt_mask用于Content Stream,它不是预测的目标(target),所以可以利用自己的信息。而Query Stream就必须使用attn_mask,预测第i个Token时不能利用自己的内容。
第三段
上面的output_h和output_g分别是two-stream的初始输入。
第四段
seg_embed是Segment embedding,XLNet是相对Segment编码:如果两个Token在同一个Segment则是0;否则是1。所以第二维是2。
在阅读seg_mat那段时我们首先需要熟悉Tensorflow增加维度的方法。
seg_id是[128, 8],那么seg_id[:, None]的shape呢?有的读者可能会猜测是[128, 8, 1](我一开始也是这么以为),但这是不对的。[:, None]的意思是在第二个维度增加一维,而原来的第二维(8)被往后推到第三维了,因此seg_id[:, None]的shape是[128, 1, 8],它等价于seg_id[:, None, :]。而cat_ids[None, :]是在第一个维度增加一维,因此是[1, 224, 8]。
接下来tf.equal(seg_id[:, None], cat_ids[None, :])会首先进行broadcasting:
计算的结果是:如果(i,j)的seg_id相同,则为True,否则为False。注意i的取值范围是0-127;而j是0-223。因为我们计算Attention时是用128个去attend to 224(包括96个mem)。最后在用tf.logical_not反过来:1表示在不同Segment而0表示同一个Segment。
最后表示成one-hot的方式(在同一个Segment为[1,0];不同的Segment为[0,1]),变成[128, 224, 8, 2],第四位就是one-hot。
本想一气把train函数写完,但今天太晚了就先到这里吧。
显示Disqus评论(需要科学上网,有Disqus的广告)
创建于: 2020-04-11 10:27:38
目录: default
标签: 无
这个函数的返回值是Estimator的一个函数和一个record_info对象,这个record_info对象记录读取的TFRecord的num_batch和file_names等数据;而第一个函数执行后的返回值是一个tf.data.Dataset,这是Tensorflow标准的输入方式,不了解的读者可以参考或者里的相关内容。
我们这里使用普通的float(float32)而不是压缩的bfloat16,关于bfloat16,有兴趣的读者可以参考。最终它们都是调用two_stream_loss函数。
请继续阅读。
原网址: