Pytorch之pack_padded_sequence与pad_packed_sequence函数

因为专注于NLP的学习与研究，所以经常接触文本数据，文本数据有一个不好的地方是不同的文本的长度不同，而输入到模型中一个Batch的数据要求他们的长度是相同的，这就产生了要给文本加padding，使得一个Batch中所有的文本长度变成相同的，但pad不能参与到训练，这就给广大的炼丹师带来了麻烦，但幸好Pytorch给我们提供了两个函数pack_padded_sequence与pad_packed_sequence让我们很好的解决了这个问题。

import torch
from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence

为了方便演示这两个函数的作用，我们需要一个toydata来演示一下，数据如下：batch_size=3，max_seq_len=7，padding用0表示。

sample = torch.tensor([[1,2,3,4,5,6,7],[1,2,3,4,0,0,0],[1,2,3,4,5,0,0]]).T
sample
tensor([[1, 1, 1],
        [2, 2, 2],
        [3, 3, 3],
        [4, 4, 4],
        [5, 0, 5],
        [6, 0, 0],
        [7, 0, 0]])

我们用一个list来储存三个句子的长度：

lengths = (7,4,5)

pack_padded_sequence
这个函数用通俗的话来说是将pad好了的sequence解开，恢复为原来没有pad的样子。

pack_padded_sequence(inputs, lengths, batch_first=False, enforce_sorted=True)

输入：

inputs (Tensor): pad好了的sequence集，数据的size要求为[T, B, *]，其中T表示的是所有句子中最长的句子的长度，B为Batch的大小，要求inputs至少是两维；
lengths (Tensor or list or tuple or …): 储存了每个句子没有pad之前的长度的集合
batch_first: 如果是True的话，就把B放在T前
enforce_sorted: 如果是True的话就说明输入的inputs已经按照句子长度递减排好了，False的话就要在函数里排。

pack_sample = pack_padded_sequence(sample, lengths, enforce_sorted=False)
print(f"pack_sample.data = {pack_sample.data}")
print(f"pack_sample.batch_sizes = {pack_sample.batch_sizes}")
print(f"pack_sample.sorted_indices = {pack_sample.sorted_indices}")
print(f"pack_sample.unsorted_indices = {pack_sample.unsorted_indices}")
pack_sample.data = tensor([1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4, 5, 5, 6, 7])
pack_sample.batch_sizes = tensor([3, 3, 3, 3, 2, 1, 1])
pack_sample.sorted_indices = tensor([0, 2, 1])
pack_sample.unsorted_indices = tensor([0, 2, 1])

我们可以直接把pack_sample输入到RNN中进行运算了。

pad_packed_sequence
这个函数就是上一个函数的逆过程。

pad_pack_sample, lengths = pad_packed_sequence(pack_sample)
pad_pack_sample
tensor([[1, 1, 1],
        [2, 2, 2],
        [3, 3, 3],
        [4, 4, 4],
        [5, 0, 5],
        [6, 0, 0],
        [7, 0, 0]])

lengths

tensor([7, 4, 5])

再来一个例子
1、准备例子数据

import torch

batch_size = 3   # 这个batch有3个序列
max_len = 6       # 最长序列的长度是6
embedding_size = 8 # 嵌入向量大小8
hidden_size = 16   # 隐藏向量大小16
vocab_size = 20    # 词汇表大小20

input_seq = [[3, 5, 12, 7, 2, ], [4, 11, 14, ], [18, 7, 3, 8, 5, 4]]
lengths = [5, 3, 6]   # batch中每个seq的有效长度。
# embedding
embedding = torch.nn.Embedding(vocab_size, embedding_size, padding_idx=0)
# GRU的RNN循环神经网络
gru = torch.nn.GRU(embedding_size, hidden_size)
2、排序数据

```python
'由大到小排序'
input_seq = sorted(input_seq, key = lambda tp: len(tp), reverse=True)
lengths = sorted(lengths, key = lambda tp: tp, reverse=True)
'''
outputs:
input_seq: [[18, 7, 3, 8, 5, 4], [3, 5, 12, 7, 2], [4, 11, 14]]
lengths : [6, 5, 3]
'''

3、填充数据

PAD_token = 0 # 填充下标是0
def pad_seq(seq, seq_len, max_length):
	seq = seq
	seq += [PAD_token for _ in range(max_length - seq_len)]
	return seq

pad_seqs = []  # 填充后的数据
for i,j in zip(input_seq, lengths):
	pad_seqs.append(pad_seq(i, len_i, max_len))
'''
填充后数据
pad_seqs : [[18, 7, 3, 8, 5, 4], [3, 5, 12, 7, 2, 0, 0], [4, 11, 14, 0, 0, 0, 0, 0, 0]]
'''

4、使用pack和pad函数

pad_seqs = torch.tensor(pad_seqs)
embeded = embedding(pad_seqs)

# 压缩，设置batch_first为true
pack = torch.nn.utils.rnn.pack_padded_sequence(embeded, lengths, batch_first=True)
'这里如果不写batch_first,你的数据必须是[s,b,e]，不然会报错lenghth错误'

#　利用gru循环神经网络测试结果
state = None
pade_outputs, _ = gru(pack, state)
# 设置batch_first为true;你可以不设置为true,为false时候只影响结构不影响结果
pade_outputs, others = torch.nn.utils.rnn.pad_packed_sequence(pade_outputs, batch_first=True)

# 查看输出的元祖
print(pade_outputs.shape) 'torch.Size([3, 6, 16])'
print(others) 'tensor([6, 5, 3])'