NLP中各框架对变长序列的处理全解

©PaperWeekly 原创 · 作者｜海晨威

学校｜同济大学硕士生

研究方向｜自然语言处理

在 NLP 中，文本数据大都是变长的，为了能够做 batch 的训练，需要 padding 到相同的长度，并在实际训练中忽略 padding 部分的影响。

在不同的深度学习框架中，对变长序列的处理，本质思想都是一致的，但具体的实现方式有较大差异，下面针对 Pytorch、Keras 和 TensorFlow 三大框架，以 LSTM 模型为例，说明各框架对 NLP 中变长序列的处理方式和注意事项。

PyTorch

在 pytorch 中，是用的 torch.nn.utils.rnn 中的 pack_padded_sequence 和 pad_packed_sequence 来处理变长序列，前者可以理解为对 padded 后的 sequence 做 pack（打包/压紧），也就是去掉 padding 位，但会记录每个样本的有效长度信息；后者是逆操作，对 packed 后的 sequence 做 pad，恢复到相同的长度。

def pack_padded_sequence(input, lengths, batch_first=False, enforce_sorted=True):
    ...
    if enforce_sorted:
        sorted_indices = None
    else:
        lengths, sorted_indices = torch.sort(lengths, descending=True)
        sorted_indices = sorted_indices.to(input.device)
        batch_dim = 0 if batch_first else 1
        input = input.index_select(batch_dim, sorted_indices)
    ...

不过在使用过程中，要格外注意 pack_padded_sequence 的 enforce_sorted 参数和 pad_packed_sequence 的 total_length 参数。

1.1 pack_padded_sequence

下面是 pack_padded_sequence 函数的部分 Pytorch 源码，input 就是输入的一个 batch 的 tensor，lengths 是这个 batch 中每个样本的有效长度。

在 pack_padded_sequence 处理之后，会得到一个 PackedSequence 的数据，其除了记录 Tensor data 之外，还会记录 batch_sizes, sorted_indices 和 unsorted_indices，其中 batch_sizes 是将输入按照有效长度排序之后，每个时间步对应的 batch 大小，后面会有例子；sorted_indices 就是对输入 lengths 排序后的索引，unsorted_indices 是用来将排序数据恢复到原始顺序的索引。

在 pack_padded_sequence 中，enforce_sorted 默认设置为 True，也就是说输入的 batch 数据要事先按照长度排序，才能输入，实际上，更简单的方式是，将其设置为 False，从上面的代码中也可以看出，Pytorch 会自动给我们做排序。

注：torch1.1 及之后才有 enforce_sorted 参数，因此 torch1.1 之后才有自动排序功能。

一个简单的例子：

# input_tensor shape：batch_size=2,time_step=3,dim=1
input_tensor = torch.FloatTensor([[4, 0, 0], [5, 6, 0]]).resize_(2, 3, 1)
seq_lens = torch.IntTensor([1, 2])
x_packed = nn_utils.rnn.pack_padded_sequence(input_tensor, seq_lens, batch_first=True, enforce_sorted=False)

输出的 x_packed 为：

PackedSequence(data=tensor([[5.],
        [4.],
        [6.]]), batch_sizes=tensor([2, 1]), sorted_indices=tensor([1, 0]), unsorted_indices=tensor([1, 0]))

在上面的例子中，首先，经过 pack_padded_sequence 内部按有效长度逆序排列之后，输入数据会变成：

[[5, 6, 0],
[4, 0, 0]]

PackedSequence 中的 data 是按照 time_step 这个维度，也就是按列来记录数据的，但是不包括 padding 位

该图仅作为理解参考，图片来自：

https://www.cnblogs.com/lindaxin/p/8052043.html

batch_sizes 记录的每列有几个数据是有效的，也就是每列有效的 batch_size 长度，但是不包括为 0 的长度，因此上面例子中，x_packed 的 batch_sizes=tensor([2, 1])，因此，每个 time_step 只需要传入对应 batch_size 个数据即可，可以减少计算量。

要注意的是，batch_sizes 这个 tensor 的长度是 2，而 input_tensor 的 time_step 是 3，因为 batch_sizes 不包含都是 padding 的时间步，也就是上面的第三列，因此后面的 pad_packed_sequence 要注意设置 total_length 参数。

1.2 pad_packed_sequence

下面是 pad_packed_sequence 函数的部分 Pytorch 源码，输入 sequence 是 PackedSequence 型数据。pad_packed_sequence 实际上就是做一个 padding 操作和根据索引恢复数据顺序操作。

def pad_packed_sequence(sequence, batch_first=False, padding_value=0.0, total_length=None):
  max_seq_length = sequence.batch_sizes.size(0)
    if total_length is not None:
        max_seq_length = total_length
    ...

这里要注意的一个参数是 total_length，它是 sequence 需要去被 padding 的长度，我们期望的一般都是 padding 到和输入序列一样的 time_step 长度，但是PackedSequence 型数据并没有记录这个数据，因此它用的是 sequence.batch_sizes.size(0)，也就是 batch_sizes 这个 tensor 的长度。

上面已经提到，batch_sizes 不包含都是 padding 的时间步，这样，如果整个 batch 中的每条记录有都做padding，那 batch_sizes 这个 tensor 的长度就会小于 time_step ，就像上面代码中的例子。

这时如果没有设置 total_length，pad_packed_sequence 就不会 padding 到我们想要的长度。

可能你在实际使用时，不设置 total_length 参数也没有出现问题，那大概率是因为你的每个 batch 中，都有至少一条记录没有 padding 位，也就是它的每一步都是有效位，那 sequence.batch_sizes.size(0) 就等于 time_step。

1.3 使用方式

为了方便使用，这里将 pack_padded_sequence，LSTM 和 pad_packed_sequence 做了一个封装，参数和原始 LSTM 一样，唯一的区别是使用中要输入 seq_lens 数据。

class MaskedLSTM(Module):
    def __init__(self, input_size, hidden_size, num_layers=1, bias=True, batch_first=False, dropout=0., bidirectional=False):
        super(MaskedLSTM, self).__init__()
        self.batch_first = batch_first
        self.lstm = LSTM(input_size, hidden_size, num_layers=num_layers, bias=bias,
             batch_first=batch_first, dropout=dropout, bidirectional=bidirectional)

    def forward(self, input_tensor, seq_lens):
        # input_tensor shape: batch_size*time_step*dim , seq_lens: (batch_size,)  when batch_first = True
        total_length = input_tensor.size(1) if self.batch_first else input_tensor.size(0)
        x_packed = pack_padded_sequence(input_tensor, seq_lens, batch_first=self.batch_first, enforce_sorted=False)
        y_lstm, hidden = self.lstm(x_packed)
        y_padded, length = pad_packed_sequence(y_lstm, batch_first=self.batch_first, total_length=total_length)
        return y_padded, hidden

小总结：

使用 pack_padded_sequence 和 pad_packed_sequence 之后，LSTM 输出对应的 padding 位是全 0 的，隐藏层输出 (h_n,c_n) 都是不受 padding 影响的，都是 padding 前最后一个有效位的输出，而且对单向/双向 LSTM 都是没有影响的，因为 padding 位不参与运算，即减少了不必要的计算，又避免了 padding 位对输出的影响。

Keras

在 keras 中，自带有 Masking 层，简单方便，使用了一个 mask 操作则可以贯穿后面的整个模型，实际的过程是把一个布尔型的 mask 矩阵一直往下游传递下去，当然这个矩阵的维度会根据当前层的维度情况重新调整，以使其能在下游层中被使用。

确实方便，但也因此丢失了灵活性，如果使用了 mask，则后面层都要支持mask，否则会报异常，这对于一些不支持 mask 的层，例如 Flatten、AveragePooling1D 等等，并不是很友好。

keras 中对于变长序列的处理，一般使用 Masking 层，如果需要用到 Embedding 层，那可以直接在 Embedding 中设置 mask_zero=True，就不需要再加 Masking 层了，但本质上都是建了布尔型的 mask 矩阵并往下游传递下去。

下面是 Masking 和 Embedding 层的定义：

Masking(mask_value=0.,input_shape=(time_step,feature_size))
Embedding(input_dim, output_dim, mask_zero=False, input_length=None)

下面是 Embedding 层中的 mask 计算函数，如果 mask_zero 设置为 True，那这里会计算 mask 矩阵并往后传递，如果要继续深入其传递的机制，建议看 keras 源码，也可以参考一下这个：keras 源码分析之 Layer [1]

# Embedding 层中的mask计算函数
def compute_mask(self, inputs, mask=None):
    if not self.mask_zero:
        return None
    output_mask = K.not_equal(inputs, 0)
    return output_mask

不过要注意的一点是，mask_zero 设置为 True，输入通过 Embedding 后，padding 位所对应的向量并不是全 0，仍然是一个随机的向量，和 mask_zero 的值没有关系，mask_zero 只是影响是否计算 mask 矩阵。但是有了 mask 矩阵之后，padding 位都不会被计算，因此，其对应向量的值并不重要。

2.1 使用方式

input = keras.layers.Input((time_step,feature_size))
mask = keras.layers.Masking(mask_value=0, input_shape=(time_step,feature_size))(input)
lstm_output = keras.layers.LSTM(hidden_size, return_sequences=True)(mask)
model = Model(input, lstm_output)

或：

input = keras.layers.Input((time_step,))
embed = keras.layers.Embedding(vocab_size, embedding_size, mask_zero=True)(input)
lstm_output = keras.layers.LSTM(hidden_size, return_sequences=True)(emd)
model = Model(input, lstm_output)

keras 模型中，Masking 之后的层，只要支持 mask，都不用再手动创建 mask 了，当然，如果是自己定义的层，要支持 mask，需要设置 supports_masking=True，并实现自己的 compute_mask 函数。

要注意的是，和 pytorch、TF 有些不一样的地方，对于有了 Masking 层之后的 LSTM，padding 位的输出不会是全 0，而是最后一位有效位的输出，也就是 padding 位输出都复制了最后有效位的输出。

Embedding 层和 Masking 层都有 mask 功能，但与 Masking 层不同的是，Embedding 它只能过滤 0，不能指定其他字符。

TensorFlow

在 TF (tf 1.x) 中是通过 dynamic_rnn 来实现变长序列的处理，它和 pytorch 的 pack_padded_sequence 一样，也有 sequence_length 参数，但它相对比 pytorch 更方便，不用手动去 pack 和 pad，只要传递 sequence_length 参数，其他都由 dynamic_rnn 来完成。

但是 TF 中 dynamic_rnn 计算的循环次数仍然是 time_steps 次，并没有带来计算效率上的提升。sequence length 的作用只是在每个序列达到它的实际长度后，把后面时间步的输出全部置成零、状态全部置成实际长度那个时刻的状态。

这一点可以参考：

https://www.zhihu.com/question/52200883

3.1 使用方式

# 静态图定义部分
basic_cell = tf.nn.rnn_cell.LSTMCell(hidden_size)
X = tf.placeholder(tf.float32, shape=[None, time_step, dim])
seq_length = tf.placeholder(tf.int32, [None])
outputs, states = tf.nn.dynamic_rnn(basic_cell, X, dtype=tf.float32, sequence_length=seq_length)

以上是从应用和代码的角度，介绍了 Pytorch、Keras 和 TensorFlow 三大框架对变长数据的处理和使用方式，但这不仅仅适用于 NLP 领域，只是在 NLP 中变长数据更为常见，希望能帮助你在工程实践中更好地去处理变长的数据。