TFSEQ Part II: 序列模型的实现细节

TFSEQ Part II: 序列模型的实现细节

本文作者：追一科技算法工程师 Tony

1. 前言

TFSEQ 这个系列总结了笔者在使用 tensorflow 进行自然语言处理的一些实践经验和思考。计划写三篇文章：

1. 分布式训练的方案和效率对比
2. 序列模型的实现细节
3. Batch size大小，优化和泛化

此为第二篇。

序列模型组件如 RNN 和 Attention 在自然语言处理中有广泛的应用。但由于序列长度不一且变化范围较大，为保证效率和稳定性，有许多实现上的细节需要考虑。同样的理论在不同实现下效果往往会有神秘的差异，甚至会出现不能收敛的情况。本文总结了一些用 Tensorflow 实现序列模型的一些做法，并分析了效率和精度上的权衡。

本文假设读者已经有深度学习在自然语言处理应用上的基本知识，并用 Tensorflow 实现过一些序列模型。为了避免翻译带来的歧义，部分术语会直接使用英文表述(使用中文的话会在括号里加上英文术语)，所以中英混杂的文风难以避免。为了讨论方便，以下先做一些术语的规定。

min-batch SGD 是一种迭代式优化(iterative optimization)的算法，每一次迭代都包括以下三个步骤：

1. 读取 mini-batch，使用模型进行前馈计算(feedforward or forward)
2. 计算 loss，并利用 loss 的值进行反向传播(backpropagation or backprop)，得到各个参数的梯度(gradient)
3. 根据算出的梯度，利用选定的优化算法(Tensorflow 中称为 Optimizer)，如标准 SGD 或者更加流行的 Adam 对参数进行更新。

数据集通常会先进行随机打乱(shuffle)，然后再将整个数据集分成固定大小(batch size)的mini-batch。这是为了保证模型见到的数据是独立同分布的(i.i.d., independent and identically distributed)，从而达到无偏的梯度估计。整个数据集全部输入并训练一遍称为一个 epoch。通常每个 epoch 开始之前都会将数据集 shuffle 一遍。

2. 数据预处理的加速

数据预处理的过程即从原始数据到将切分好的 mini-batch 载入模型的整个过程。tensorflow 的文档将整个流程划分为 Extract，Transform 和 Load(ETL)

A typical TensorFlow training input pipeline can be framed as an ETL process:

1. Extract: Read data from persistent storage – either local (e.g. HDD or SSD) or remote (e.g. GCS or HDFS).
2. Transform: Use CPU cores to parse and perform preprocessing operations on the data such as image decompression, data augmentation transformations (such as random crop, flips, and color distortions), shuffling, and batching.
3. Load: Load the transformed data onto the accelerator device(s) (for example, GPU(s) or TPU(s)) that execute the machine learning model.

2.1 Placeholder 模式

tensorflow 最初推荐使用的是 tf.placeholder 模式。这个模式的流程如下：

1. Extract：先将整个数据集读入内存
2. Transform：在Python中将数据转化成 numpy 的矩阵
3. Load：使用 tf.placeholder 将 numpy 数据拷贝到模型中，再由模型拷贝到 GPU 中。

[图注]tf.placeholder 模式的图示，虚线框值指的是计算操作，实线框指的是存储。曲线数指系统线程数。所有步骤共享同一线程。

这种模式虽然简单，但缺点也很多：

1. 当数据集很大的时候，可能无法将整个数据集载入内存。
2. python 代码的运行速度较慢，多线程由于 python GIL 的存在无法有效并行
3. 多了不必要的拷贝

最明显的缺点是，在整个处理流程中 GPU 一直是闲置的。

所以 tensorflow 如今主推的是 tf.data api，tf.placeholder的介绍已经很难在文档中找到了(这也体现了 Tensorflow 文档对新老用户感情的玩弄)。但在实际线上服务的时候，tf.placeholder 仍然是把数据载入模型的主要方式。

2.2 Input pipeline

input pipeline 在 tensorflow 早期就已经存在，真正做到易用还是从 tensorflow 1.2 发布 tf.contrib.data 接口开始，其前身是杂乱无章的 QueueRunner 和 Queue。

为了解决 tf.placeholder 的效率问题，tensorflow 采用了 cache 和 pipeline 的思想。

1. Extract：将数据集顺序读入一个固定长度的 queue 中。如果 queue 中的数据达到其容量上线，则阻塞直到下游任务读取数据空出容量。
2. Transform：从 extract 的 queue 中读取(dequeue)一定量的数据，开启多条线程并行处理数据，将得到的 mini-batch 放在一个固定长度的 queue 中。
3. Load：从 transform 后的 queue 中读取(dequeue) 一个 mini-batch，传到 GPU。新版本的tensorflow 还提供了 tf.contrib.data.prefetch_to_device 接口，支持将上一步的queue 放在 GPU 内存里，从而进一步减少等待时间。

通过维护一个固定长度的 queue 作为 cache 来解决内存占用的问题。三个步骤通过 cache 独立开来，有独立的线程去维护，需要调整参数保证各个 queue 不会排空，从而避免阻塞。为了加快处理速度，transform 这一步还加入了多线程支持。

下图总结了各个步骤的核心函数。曲线数量指并行线程数(系统线程而不是 python 线程)。列出的函数还有其他的参数，图中写出的参数是核心参数。虚线框值指的是计算操作，实线框指的是存储。

[图注]tf.data 模式的图示，虚线框值指的是计算操作，实线框指的是存储。曲线数指系统线程数。所有步骤都有各自的线程。

对于文本数据而言，一个典型的 dataset 如下

def build_dataset(filenames,
                  parse_func,
                  filter_func,
                  pad_id,
                  batch_size,
                  buffer_size,
                  shuffle=False,
                  repeat=False):
    dataset = tf.data.TextLineDataset(filenames,
                                      buffer_size=buffer_size)
    if shuffle:
        dataset = dataset.shuffle(buffer_size)
    if repeat:
        # repeat forever
        dataset = dataset.repeat()
    # filter invalid lines out
    dataset = dataset.filter(lambda line: filter_func(line))
    dataset = dataset.map(lambda line: parse_func(line),
                          num_parallel_calls=2)
    # suppose element is (word_ids, length)
    dataset = dataset.padded_batch(batch_size,
                                   padded_shapes=(
                                       tf.TensorShape([None]),
                                       tf.TensorShape([])),
                                   padding_values=(pad_id, 0))
    dataset = dataset.prefetch(8)
    return dataset

2.3 Data pipeline 的优化

理解上述过程后，对 buffer_size 和 num_parallel_calls 这两个参数的调试目标就是在保证 cache 不被排空的基础上，尽量少地占用计算资源(如果你还有别的任务要跑的话)。tf.data.Dataset.prefetch 通常只放在最后一步，其 buffer_size 一般设成 8 左右。对于不复杂的 Transform，tf.data.Dataset.map 的 num_parallel_calls 一般设成 1~2。

pipeline 的运行和模型训练是独立的，所以只需要保证处理数据的速度比训练的速度快即可，而这在文本数据处理中很容易做到。由于 tensorflow 对文本处理的支持不好，建议将比较复杂的文本处理先离线做好存成数据文件，再用 tensorflow 读取。

需要注意的是 tf.data.Dataset.shuffle 的使用。因为 shuffle 是在 Queue 上做的，当 buffer_size 比较小时，对数据集的 shuffle 会不够充分。实际使用时建议先将硬盘中的数据全局 shuffle 一遍： shuf data.txt > shuf_data.txt，并将 buffer_size 设大一点。

2.4 Data pipeline 的技巧

在你的训练流程中，如果 GPU 的计算需要依赖没有 GPU 实现的操作，而这部分操作又不会很重的话，可以把它放在 input pipeline 中，充分利用 tf.data.Dataset.prefetch 的 cache 功能。这里举两个例子：

1. 在实现 GPU 版 word2vec 时，可以把 negative sampling 部分放在 input pipeline 中。
2. 在实现多卡训练的时候，如果只维护一条 input pipeline，可以将大 batch 切分的处理放在 input pipeline 里做。个人发现这比这个讨论提到的方案都要好一些。

def split_batches(num_splits, batches):
    batch_size = tf.shape(batches[0])[0]
    # evenly distributed sizes
    divisible_sizes = tf.fill([num_splits], 
    													tf.floor_div(batch_size, num_splits))
    remainder_sizes = tf.sequence_mask(tf.mod(batch_size, num_splits),
                                       maxlen=num_splits,
                                       dtype=tf.int32)
    frag_sizes = divisible_sizes + remainder_sizes
    batch_frags_list = []
    for batch in batches:
        batch_frags = tf.split(batch, frag_sizes, axis=0)
        batch_frags_list.append(batch_frags)

    frag_batches_list = zip(*batch_frags_list)
    # fix corner case
    for i, frag_batches in enumerate(frag_batches_list):
        if len(frag_batches) == 1:
            frag_batches_list[i] = frag_batches[0]
    return frag_batches_list

dataset = dataset.map(lambda seqs_batch, lengths_batch:
                      split_batches(num_splits,
                                    [seqs_batch, lengths_batch]),
                      num_parallel_calls=4)

更多的建议(如fused transformation)可以参照官方文档。

3. 变长序列的处理

3.1 处理 mini-batch 中的序列长度变化：Padding 和 Bucketing

由于数据集中的序列的长度不一，我们需要对 mini-batch 中的短句进行补长(padding)的操作：可用指定的符号(token)拼到短句末位，将短句补到 mini-batch 里面最长句子的长度。这是因为模型的大部分操作都是基于矩阵运算。补长用的符号对应的特征通常会在后续的计算开始前置零。一般补长操作都是在做完词表映射后做的。

当 mini-batch 中的句长范围过大的时候(如1~100)，GPU 大部分时间都在处理最长的句子，短句有效的部分很快就处理完了。为了更充分利用GPU的并行计算资源，我们可以控制 mini-batch 的句长范围。

Bucketing 的想法就是在构建 mini-batch 之前，先将序列按长度分组(bucket)，然后再从每个 bucket 里面随机选取序列构建 mini-batch。按长度分组的操作可以通过简单的 hash 实现。

Bucketing + padding 可以借由 tf.data 接口实现：

def bucket_pad_batch(dataset, pad_id):
	# dataset element is (seq, length). seq is an integer tuple.
    def key_func(seq, length):
        bucket_id = length // bucket_width
        return tf.to_int64(bucket_id)

    def reduce_func(unused_key, windowed_data):
        return windowed_data.padded_batch(
                batch_size,
                padded_shapes=(
                    tf.TensorShape([None]),
                    tf.TensorShape([]), # 表示不做 padding
                ),
                padding_values=(pad_id, 
                                0, # 因为不做padding，这只是占位符
                               ))
    dataset = dataset.apply(tf.contrib.data.group_by_window(
            key_func=key_func,
            reduce_func=reduce_func,
            window_size=batch_size))
    return dataset

当数据集的句长分布范围较大的时候，bucketing 可以显著提升训练速度。

但 bucketing 分组的操作会妨碍 shuffle 的效果，使训练数据偏离 i.i.d.(Independent and identically distributed) 的假设，使梯度估计有偏，在训练过程中使 loss 的抖动加剧，模型的收敛性变差。实际使用中需要调整 bucket 的数量，对训练速度和收敛性能进行 trade-off。在这篇水文中有这个图可以参考:

[图注]Bucket 数量对训练速度和收敛速度的影响。图来自Accelerating recurrent neural network training using sequence bucketing and multi-GPU data parallelization

3.2 变长序列的 Soft Attention

soft attention 的计算方法如下:

$a_{{\text{word}}_t}=\frac{\exp (s_{{\text{word}}_t})}{{\sum }_i\exp (s_{{\text{word}}_i})}$

其中 $a_{{\text{word}}_t}$ 为第 $t$ 个词的权重，$s_{{\text{word}}_t}$ 为 attention key 和第 $t$ 个词算出来的相似度。我们可以简单地调用 tf.nn.softmax 进行计算。attention 后的句子特征为

$w={\sum }_t{a}_{{\text{word}}_t}{w}_{{\text{word}}_t}$

其中 $w_{{\text{word}}_t}$ 是第 $t$ 个词的特征向量。

# 实现 1
# feat_seqs: shape=[batch_size, maxlen, feature] # 每个词的特征
# sims: shape=[batch_size, maxlen] # 每个词的 attention 相似度
# lengths: shape=[batch_size] # 句长
atns = tf.nn.softmax(sims) # [batch_size, maxlen]

atns = tf.expand_dims(atns, axis=-1) # [batch_size, maxlen, 1]
feat_seqs = tf.transpose(feat_seqs, [0,2,1]) # [batch_size, feature, maxlen]
feat = tf.matmul(feat_seqs, atns)

在 mini batch 训练过程中，大多数序列后面都会带上占位符 <pad>，如果不做处理的话，上述实现对 attention 权重的计算应为下式：

a ^ word t = exp ⁡ ( s word t ) ∑ i exp ⁡ ( s