pytorch中的pad_sequence、pack_padded_sequence和pad_packed_sequence函数

在使用pytorch训练模型的时候，一般采用batch的形式同时处理多个样本序列，而同一batch中时序信息的的长度是不同的，这样就无法传入RNN，LSTM，GRU这样的模型中进行处理。一个常用的做法是按照一个指定的长度(或者按照batch中最长的序列长度)对batch中的序列进行填充(padding)或者截断(truncate)，这样就会导致一些较短的序列中会有很多的填充符。如下图所示：

这样就会导致一个问题，就是对于进行padding过的序列，LSTM对它进行处理时处理了很多无用的字符，这样得到的句子的表征就会有误差，如下图所示:

而我们希望得到的序列的表征是处理过原有长度数据之后的表征，而不是处理过很多个没用的填充符之后的表征：

一般使用torch.nn.utils.rnn.pad_sequence对序列进行填充，填充的长度为批次中最长序列的长度，再对序列进行填充之后，就需要用torch.nn.utils.rnn.pack_padded_sequence 和torch.nn.utils.rnn.pad_packed_sequence两个函数进行处理，处理的过程如下：

1. 填充后的输入序列先经过torch.nn.utils.rnn.pack_padded_sequence的处理，然后会得到一个PackedSequence类型的对象，可以直接给RNN进行处理。
2. RNN处理PackedSequence类型的数据后，会返回一个PackedSequence类型的输出。
3. 最后使用torch.nn.utils.rnn.pad_packed_sequence函数将经过RNN后的输出数据在重新进行填充，得到正常的每个batch等长的序列。

下面看看这三个函数。

1.1 、torch.nn.utils.rnn.pad_sequence

再pytorch中该函数的形式为：

torch.nn.utils.rnn.pad_sequence(sequences, batch_first=False, padding_value=0.0)

该函数用padding_value来填充一个可变长度的张量列表。将长度较短的序列填充为和最长序列相同的长度。

参数说明：

• sequences(list[Tensor])：变长序列的列表。
• batch_frist(bool,optional)：如果为True，output形状为$B\times T\times \ast$，否则为$T\times B\times \ast$，默认情况为False。其中$B$为批次大小，$T$为填充后每个序列的长度。
• padding_value(float,optional)：填充元素的值。默认值：0。

输出：

如果 batch_first 是 False，张量的形状为$T\times B\times \ast$。否则，张量的形状为$B\times T\times \ast$。

例子1：

对三个不等长的张量进行填充。

import torch

a = torch.ones(25, 300)
b = torch.ones(22, 300)
c = torch.ones(15, 300)
#batch_first = False
torch.nn.utils.rnn.pad_sequence([a, b, c]).size()
'''
torch.Size([25, 3, 300])
'''

#batch_first = True
torch.nn.utils.rnn.pad_sequence([a, b, c],batch_first = True).size()
'''
torch.Size([3, 25, 300])
'''

例子2：

inputs = [
    torch.tensor([[1,2,3],[1,2,3],[1,2,3]]),
    torch.tensor([[1,2,3],[1,2,3],[1,2,3],[1,2,3]]),
    torch.tensor([[1,2,3],[1,2,3]])  
]
torch.nn.utils.rnn.pad_sequence(inputs,batch_first = True)
'''
tensor([[[1, 2, 3],
         [1, 2, 3],
         [1, 2, 3],
         [0, 0, 0]],

        [[1, 2, 3],
         [1, 2, 3],
         [1, 2, 3],
         [1, 2, 3]],

        [[1, 2, 3],
         [1, 2, 3],
         [0, 0, 0],
         [0, 0, 0]]])
'''

1.2、torch.nn.utils.rnn.pack_padded_sequence

压紧(pack)一个包含可变长度的填充序列的张量，在使用torch.nn.utils.rnn.pad_sequence函数进行填充的时候，产生了冗余，因此需要对其进行pack。在pack的过程中时按照列进行pack，而不是按照行进行pack。如下图过程：

Pytorch中的格式为：

torch.nn.utils.rnn.pack_padded_sequence(input, lengths, batch_first=False, enforce_sorted=True)

参数说明：

• input(Tensor)：一批量填充后的可变长度的序列。
• lenghts(Tensor or list(int))：每个批次元素的序列长度列表。如果输入为张量形式则必须在CPU上，不能在GPU上！！！一个小坑。
• batch_first(bool,optional)：如果为True，则输入的形状为$B\times T\times \ast$，我一般将其设置为True。
• enforce_sorted(bool,optional)：如果为True，则参数lenghts为按长度递减排序的序列，这样的话输入的input也需要进行排序。我一般将其设置为False。如果为False输入将被无条件地排序

输出：

一个PackedSequence对象，下面通过一个例子看看这个对象中都包含什么。

例子：

data = [torch.tensor([9]),
       torch.tensor([1,2,3,4]),
       torch.tensor([5,6])]
seq_len = [1,4,2]
#padding the sequence
pad_data = torch.nn.utils.rnn.pad_sequence(data, batch_first=True)
#pack data
pack_data = torch.nn.utils.rnn.pack_padded_sequence(pad_data, seq_len, batch_first=True, enforce_sorted=False)
pack_data
'''
PackedSequence(data=tensor([1, 5, 9, 2, 6, 3, 4]), batch_sizes=tensor([3, 2, 1, 1]), sorted_indices=tensor([1, 2, 0]), unsorted_indices=tensor([2, 0, 1]))
'''

data为pack后的数据，batch_sizes为按照列进行pack后每一列数据的个数，

sorted_indices表示PackedSequence是如何从序列中构建的。unsorted_indices表示如何以正确的顺序恢复原始序列。

1.3、torch.nn.utils.rnn.pad_packed_sequence

该函数的功能是对可变长度的序列进行填充。主要使用场景为，将pack的数据传入RNN，LSTM或者GRU这样的模型后，返回的输出还是PackedSequence类型，因此需要使用该函数对返回的结果进行填充以恢复为原来的形状。

该函数在Pytorch中的格式为：

torch.nn.utils.rnn.pad_packed_sequence(sequence, batch_first=False, padding_value=0.0, total_length=None)

参数说明：

• sequence(PackedSequence)：需要填充的数据。
• batch_first(bool,optional)：如果为True，输出形状为$B\times T\times \ast$。
• padding_value(float,optional)：填充元素的值。
• total_lenght(int,optional)：如果不是无，输出将被填充成total_lenght。

输出：

包含填充序列的张量的元组，以及包含批次中每个序列的长度列表的张量。

例子：

import torch
import torch.nn as nn
 
#GRU model    
gru = nn.GRU(input_size=1, hidden_size=1, batch_first=True)
 
data = [torch.tensor([9.]),
       torch.tensor([1.,2.,3.,4.]),
       torch.tensor([5.,6.])]
seq_len = [1,4,2]

#padding the sequence
pad_data = torch.nn.utils.rnn.pad_sequence(data, batch_first=True).float().unsqueeze(2)

print('padding data：')
print(pad_data)
#pack data
pack_data = torch.nn.utils.rnn.pack_padded_sequence(pad_data, seq_len, batch_first=True, enforce_sorted=False)

output, hidden = gru(pack_data)
print('output of the model：')
print(output)

#pack the output
output, _ = torch.nn.utils.rnn.pad_packed_sequence(sequence=output, batch_first=True)

print('pack the output：')
print(output)

输出：

padding data：
tensor([[[9.],
         [0.],
         [0.],
         [0.]],

        [[1.],
         [2.],
         [3.],
         [4.]],

        [[5.],
         [6.],
         [0.],
         [0.]]])
output of the model：
PackedSequence(data=tensor([[ 1.9511e-02],
        [-1.6523e-03],
        [-8.1122e-05],
        [ 1.9253e-02],
        [-2.5014e-03],
        [ 1.5634e-02],
        [ 1.2749e-02]], grad_fn=<CatBackward0>), batch_sizes=tensor([3, 2, 1, 1]), sorted_indices=tensor([1, 2, 0]), unsorted_indices=tensor([2, 0, 1]))
pack the output：
tensor([[[-8.1122e-05],
         [ 0.0000e+00],
         [ 0.0000e+00],
         [ 0.0000e+00]],

        [[ 1.9511e-02],
         [ 1.9253e-02],
         [ 1.5634e-02],
         [ 1.2749e-02]],

        [[-1.6523e-03],
         [-2.5014e-03],
         [ 0.0000e+00],
         [ 0.0000e+00]]], grad_fn=<IndexSelectBackward0>)