Pytorch中pack_padded_sequence和pad_packed_sequence的理解

引言

这里补充下对Pytorch中​​pack_padded_sequence​​​和​​pad_packed_sequence​​的理解。

当我们训练RNN时，如果想要进行批次化训练，就得需要截断和填充。
因为句子的长短不一，一般选择一个合适的长度来进行截断；

而填充是在句子过短时，需要以 填充字符 填充，使得该批次内所有的句子长度相同。

​​pack_padded_sequence​​做的就是压缩这些填充字符，加快RNN的计算效率。

为什么要进行压缩填充

而填充会带来两个问题：

1. 增加了计算复杂度。假设一个批次内有2个句子，长度分别为5和2。我们要保证批次内所有的句子长度相同，就需要把长度为2的句子填充为5。这样喂给RNN时，需要计算$2 \times 5 =10$次，而实际真正需要的是$5+2=7$次。
2. 得到的结果可能不准确。我们知道RNN取的是最后一个时间步的隐藏状态做为输出，虽然一般是以0填充，权重乘以零不会影响最终的输出，但在Pytorch中还有偏差$b$，如果$b \neq 0$，还是会影响到最后的输出。当然这个问题不大。主要是第1个问题。 毕竟批次大小很大的时候影响还是不小的。

所以Pytorch提供了​​pack_padded_sequence​​方法来压缩填充字符。

如何压缩

假设一个批次有6个句子，我们将这些句子填充后如下所示。

这里按句子长度逆序排序，pads就是填充。不同的颜色代表不同时间步的单词。

下面我们看​​pack_padded_sequence​​是如何压缩的。如下图所示：

​​pack_padded_sequence​​根据时间步拉平了上面排序后的句子，在每个时间步维护一个数值，代表当前时间步内有多少个批次数据。比如上图右边黄色区域的时间步内，只有3个输入是有效的，其他都是填充。因此说该时间步内的批次数为3。

Python中​​batch_first​​不同的取值，压缩的方式有点不同，不过主要思想差不多。

该方法会返回一个​​PackedSequence​​​对象，其中包含 ​​data​​​保存拉平的数据 和 ​​batch_sizes​​​保存时间步相应的批次大小，比如上面就是​​tensor([4, 3, 3, 2, 1, 1])​​。

Pytorch的RNN(LSTM/GRU)可以接收​​PackedSequence​​​，并返回一个新的​​PackedSequence​​​。然后我们可以用​​pad_packed_sequence​​​方法把返回的​​PackedSequence​​还原成我们想要的形式。

下面以一个例子来说明。

import torch 
import torch.nn as nn

seq_batch = [torch.tensor([[1, 1],
                           [2, 2],
                           [3, 3],
                           [4, 4],
                           [5, 5]]),
             torch.tensor([[10, 10],
                           [20, 20]])]

seq_lens = [5, 2]

例子参考了下面的引用。

这里假设我们的词嵌入维度是2，并且有两个句子，第一个句子长度为5，第二个句子长度为2。

我们先进行填充，采用默认​​batch_first=False​​的方式，会返回填充后的Tensor对象。

padded_seq_batch = torch.nn.utils.rnn.pad_sequence(seq_batch, batch_first=False)

print(padded_seq_batch)
print(padded_seq_batch.size())

当​​batch_first=True​​时，填充的结果为：

我们这里主要关注默认的方式。

接下来，我们要传入LSTM中，在这之前，我们对这些填充后的Tensor进行压缩。

# 压缩
packed_seq_batch = torch.nn.utils.rnn.pack_padded_sequence(padded_seq_batch, lengths=seq_lens)

print(packed_seq_batch)
print(packed_seq_batch.data.size())

它返回一个​​PackedSequence​​对象实例，其中包括压缩后的内容，和每个时间步内批次大小。

压缩的时候，需要传入实际语句的长度，以方便后面的逆操作。

注意最新版本的Pytorch已经不需要在传入之前对句子长度进行排序。

为什么采用默认​​batch_first=False​​默认的方式，主要让结果和上面示例图的方式保持一致。

下面我们把这个​​PackedSequence​​​对象传入LSTM中，LSTM可接收​​PackedSequence​​​对象，这种情况下，只需要计算$5+2=7$次。

lstm = nn.LSTM(input_size=2, hidden_size=3)

output, (hn, cn) = lstm(packed_seq_batch.float())

当LSTM模型接收​​PackedSequence​​​对象后，返回的​​output​​​也封装在​​PackedSequence​​对象中。

此时，我们需要对输出进行解压缩，并填充回我们熟悉的形状。

解压缩

我们调用​​pad_packed_sequence​​方法进行解压缩。

padded_output, output_lens = torch.nn.utils.rnn.pad_packed_sequence(output)

这样就得到了我们熟悉的输出大小：​​(L=5,N=2,hidden_size=3)​​。

通过这种压缩、调用RNN、解压缩的方法可以批次训练，并且保持高效~！

参考

1. ​​why-do-we-pack-the-sequences-in-pytorch​​