[Pytorch]Pack the data to train variable length sequences

前言

在NLP task中常常會遇到input data長度不固定的問題，一般來說此時有兩種方法處理: 對資料padding補齊到相同長度，或是截斷超過某個長度的data。不過其實還有一種更有效率的方法，就是PackSequence，這篇文章將會針對Pytorch的PackSequence做一系列的介紹，包含概念和使用技巧。

關於PackSequence，在Pytorch的document中torch.nn.utils.rnn定義了以下的object / function:

nn.utils.rnn.PackedSequence
nn.utils.rnn.pad_sequence(): Pad a list of variable length Tensors with padding_value
nn.utils.rnn.pack_padded_sequence(): Packs a Tensor containing padded sequences of variable length.
nn.utils.rnn.pad_packed_sequence(): Pads a packed batch of variable length sequences.
nn.utils.rnn.pack_sequence(): Packs a list of variable length Tensors

主要重點在於pad_sequence(), pack_padded_sequence(), pad_packed_sequence()，而pack_sequence()其實就是pad_sequence()+pack_padded_sequence()，最後再來談。

首先，讓我們先把該import的都先寫好：

import torch
from torch import nn
import torch.utils.data as data
from torch.utils.data import DataLoader
from torch.utils.data import Dataset
import torch.nn.utils.rnn as rnn_utils
from pprint import pprint

pad_sequence

pad_sequence用來將list of tensor用0補齊，這用在你的資料輸入長度不一致的時候，由於模型輸入要求長度固定所以必須要將長度變成一樣長。

在NLP相關的task最常看到這種現象，舉例來說: 當你的輸入是句子，而每句長度都不同的時候

假設我們的資料沒有做embedding，也就是說句子裡每一個字都只是一個數字

train_x = [torch.tensor([1, 1, 1, 1, 1, 1]),
           torch.tensor([2, 2, 2, 2]),
           torch.tensor([3, 3, 3, 3]),
           torch.tensor([4, 4, 4]),
           torch.tensor([5])]
pprint(train_x)

# pad the list of sequences
pad_train_x = rnn_utils.pad_sequence(train_x, batch_first=True)
pprint(pad_train_x)

進行padding後，pad_train_x如下，將長度小於最大長度的都用0補齊了:

tensor([[1, 1, 1, 1, 1, 1],
        [2, 2, 2, 2, 0, 0],
        [3, 3, 3, 3, 0, 0],
        [4, 4, 4, 0, 0, 0],
        [5, 0, 0, 0, 0, 0]])

如果句子中的每個字都經過embedding了? 比方說現在句子中的每個字都是一個4維的向量時:

train_x1 = [torch.ones(3, 4),
            torch.ones(1, 4),
            torch.ones(5, 4)]
pprint(train_x1)

# pad the list of sequences
pad_train_x1 = rnn_utils.pad_sequence(train_x1, batch_first=True)
pprint(pad_train_x1)

padding後得到的pad_train_x1會是:

tensor([[[1., 1., 1., 1.],
         [1., 1., 1., 1.],
         [1., 1., 1., 1.],
         [0., 0., 0., 0.],
         [0., 0., 0., 0.]],

        [[1., 1., 1., 1.],
         [0., 0., 0., 0.],
         [0., 0., 0., 0.],
         [0., 0., 0., 0.],
         [0., 0., 0., 0.]],

        [[1., 1., 1., 1.],
         [1., 1., 1., 1.],
         [1., 1., 1., 1.],
         [1., 1., 1., 1.],
         [1., 1., 1., 1.]]])

關於pad_sequence()內的參數batch_first會造成什麼影響:

batch_first=True時，return的維度是: [Batch, Sequence, Features]
batch_first=False時，return的維度是: [Sequence, Batch, Features]

通常我們會設置batch_first=True，因為比較符合一般我們在思考的格式

不過這個格式餵入RNN相關模型時其實並不利於平行化計算(但其實Pytorch內部又會把他轉回batch_first=False，所以用我們比較好理解的表達形式就好)，在最後面的時候再來補充說明

對於DataLoader，DataLoader無法處理不定長度的輸入，來看一下長度不相同的時候DataLoader會花生什麼4:

train_x = [torch.tensor([1, 1, 1, 1, 1, 1]),
           torch.tensor([2, 2]),
           torch.tensor([3, 3, 3, 3]),
           torch.tensor([4, 4, 4]),
           torch.tensor([5])]
pprint(train_x)

class Mydataset(Dataset):
    def __init__(self, data):
        self.data = data
    
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        return self.data[idx]

train_loader = DataLoader(Mydataset(train_x), batch_size=2)
batch_x = iter(train_loader).next()
pprint(batch_x)

恭喜！你會得到下面的Error msg:

RuntimeError: invalid argument 0: Sizes of tensors must match except in dimension 0. Got 6 and 2 in dimension 1 at /tmp/pip-req-build-4baxydiv/aten/src/TH/generic/THTensor.cpp:689

所以要嘛在前處理就先padding好，不然就是透過DataLoader的collate_fn參數來處理從Dataset return的data:

在這邊的collate_fn中對資料先做了降序排列，原因後面會說到

def collate_fn(data):
    # sort data with descending order
    data.sort(key=lambda x: len(x), reverse=True)
    data = rnn_utils.pad_sequence(data, batch_first=True)
    return data

train_loader = DataLoader(Mydataset(train_x), batch_size=2, collate_fn=collate_fn)
batch_x = iter(train_loader).next()
pprint(batch_x)

tensor([[1, 1, 1, 1, 1, 1],
        [2, 2, 0, 0, 0, 0]])

pack_padded_sequence

為何需要pack?

在上一節，使用了pad_sequence()將長度變成一致後，其實就可以餵入RNN系列的model了。

但我們來考慮一下一組有padding的句子在RNN model中是怎麼樣訓練的，假設input sequence如下:

['a', 'b', <PAD>, <PAD>, <PAD>]

由於要與其他sequence長度一致，假設我們將這個句子補齊到長度為5，並用<PAD>代表padding

RNN(h1) -> RNN(h2) -> RNN(h3) -> RNN(h4) -> ...
 ^          ^           ^          ^  
'a'        'b'        <PAD>      <PAD>

我們通常都拿最後一個hidden state作為模型輸出，而此時的輸出是經過了許多padding sybmol的，理想上我們要拿的hidden state應該是輸入’b’當下的output(也就是h2)。

當然你可以想辦法拿出h2作為模型的output
或是說可以預期模型學得好的時候，應該學到padding symbol是無意義的，也就是說理論上h4應該要跟h2差不多才對

另外，一個很重要的議題是: 在雙向的RNN系列模型中，如果不用pack則會很麻煩(因為你需要handle兩個padding sequence: 一個往前padding + 一個往後padding)，所以雙向的RNN模型都會使用pack sequence來實作。

現在，對於上述的Issue，我們有更好的一個方式來處理: 透過pack sequence來避免訓練多餘的padding symbol

pack a sequence

應該怎麼做呢，這又扯到了model在餵入資料的時候是怎麼吃data的，用上面的example data來說明(這裡先不考慮batch_first的影響，單純先假設每一列都是一筆sequence)

tensor([[1, 1, 1, 1, 1, 1],
        [2, 2, 2, 2, 0, 0],
        [3, 3, 3, 3, 0, 0],
        [4, 4, 4, 0, 0, 0],
        [5, 0, 0, 0, 0, 0]])

RNN系列的model每一次的循環是先餵每一個sequence內第一個timestamp的data，也就是

[[1],
 [2],
 [3],
 [4],
 [5]]

然後產生下一個時間的hidden state，再計算每一個sequence內下一個timestamp的hidden state

[[1],
 [2],
 [3],
 [4],
 [0]]

依此類推，你會發現每個循環都是餵入相同筆數的data，然後產生下一個hidden state。但以第二個timestamp來看，最後一筆的[0]明明只是個padding，根本不用計算的，因此造成了額外的計算資源浪費。

所以pack在做的事情就是讓每一次的筆數省掉計算多餘的padding，你可以想像成在原本batch下我們又透過一個mini-batch來紀錄沒有padding symbol下要餵的長度，為此透過一個Class(PackedSequence)來封裝我們所需要的資訊:

PackedSequence(
  data=tensor([1, 2, 3, 4, 5, 1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 1, 1]),
  batch_sizes=tensor([5, 4, 4, 3, 1, 1])
)

PackedSequence的data是原始的資料去除掉padding後串接起來的一維data，而batch_sizes是說在原本batch下的資料透過pack轉換後，每一次的mini-batch長度(下面除非特地提到參數，不然統一mini-batch代替pack裡面的batch_sizes以免搞混)

第一個batch_sizes參數是5，代表第一個mini-batch應該從data中取出5筆data([1, 2, 3, 4, 5])來作為第一個timestamp的input
第二個batch_sizes參數的值是4，代表第二個mini-batch應該從data中取出4筆data([1, 2, 3, 4])來作為第二個timestamp的input

如此，我們的model就可以接受長度不同的input了，並且還不用考慮padding。並且這邊的mini-batch是指在同一個batch下的所有sequence去進行pack的結果，並不會影響其他batch，所以仍然可以用GPU平行計算

舉個例子，batch size=5，那可以想像成訓練的時候每張顯卡都拿不同的5筆sequences資料
然後每個GPU在各自對自己的這5比去做pack sequence取得不包含padding下的mini-batch

最後，我們仔細看一下PackedSequence這一個Class，在Pytorch doc中提到:

Holds the data and list of batch_sizes of a packed sequence.

All RNN modules accept packed sequences as inputs.

也就是說這一個類別他同時紀錄了資料和packed sequence的batch size
對於所有的RNN model，input都是可以接受packed seq的
文件中也提到，這一個Class不應該被手動產生，而應該透過對應的function來生成(pack_padded_sequence)

write a code for packing sequence

懂了pack的概念後，來看一下怎麼寫，pack_padded_sequence()顧名思義，是對有padding的sequence來進行pack這個操作，所以要先做pad_sequence()

並且要注意sequence必須是降序排列，為了避免在串接中有padding被夾在中間的原因
- 所以前面在寫collate_fn()的時候才會先對sequences做排序
pack_padded_sequence()同樣具有batch_first的參數，所以也要跟pad_sequence()的設置相同

還有一個參數lengths，用來告訴每一個sequence下真正的長度(不包含padding)是多少

train_x = [torch.tensor([1, 1, 1, 1, 1, 1]),
           torch.tensor([2, 2, 2, 2]),
           torch.tensor([3, 3, 3, 3]),
           torch.tensor([4, 4, 4]),
           torch.tensor([5])]
pprint(train_x)
# [tensor([1, 1, 1, 1, 1, 1]),
#  tensor([2, 2, 2, 2]),
#  tensor([3, 3, 3, 3]),
#  tensor([4, 4, 4]),
#  tensor([5])]

rnn_utils.pack_padded_sequence(rnn_utils.pad_sequence(train_x, batch_first=True), lengths=[6, 4, 4, 3, 1], batch_first=True)
# PackedSequence(
#   data=tensor([1, 2, 3, 4, 5, 1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 1, 1]), 
#   batch_sizes=tensor([5, 4, 4, 3, 1, 1]), sorted_indices=None, unsorted_indices=None)

對於data loader，由於長度仍然要是相同的，所以一樣會先padding，等到load進來之後才做pack，但這樣就無法得知seq的真正長度了，所以collate_fn()必須也要回傳長度的部分:

train_x = [torch.tensor([1, 1, 1, 1, 1, 1]),
           torch.tensor([2, 2]),
           torch.tensor([3, 3, 3, 3]),
           torch.tensor([4, 4, 4]),
           torch.tensor([5])]
# pprint(train_x)

class Mydataset(Dataset):
    def __init__(self, data):
        self.data = data
    
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        return self.data[idx]

def collate_fn(data):
    data.sort(key=lambda x: len(x), reverse=True)
    seq_lens = [len(seq) for seq in data]
    data = rnn_utils.pad_sequence(data, batch_first=True)
    return data, seq_lens

train_loader = DataLoader(Mydataset(train_x), batch_size=2, collate_fn=collate_fn)
batch_x, seq_lens = iter(train_loader).next()
pprint(batch_x)
# tensor([[1, 1, 1, 1, 1, 1],
#         [2, 2, 0, 0, 0, 0]])

pprint(seq_lens)
# [6, 2]

packed_batch_x = rnn_utils.pack_padded_sequence(batch_x, lengths=seq_lens, batch_first=True)
pprint(packed_batch_x)
# PackedSequence(
#   data=tensor([1, 2, 1, 2, 1, 1, 1, 1]), 
#   batch_sizes=tensor([2, 2, 1, 1, 1, 1]), sorted_indices=None, unsorted_indices=None)

pad_packed_sequence

pad_packed_sequence()就是pack_padded_sequence()的逆操作，返回原本padding的形式以及對應的length

rnn_utils.pad_packed_sequence(packed_batch_x)
# (tensor([[1, 2],
#          [1, 2],
#          [1, 0],
#          [1, 0],
#          [1, 0],
#          [1, 0]]),
#  tensor([6, 2]))

用PackedSequence寫一個RNN訓練的流程

寫一個RNN model來實際跑跑看。

首先，還記得我們說過batch_first=True時，return的維度是: [Batch, Sequence, Features]，上面最後的程式碼維度其實只是[Batch, Sequence]，因為這樣的輸出比較容易理解，不過實際上在餵入模型的時候要改一下data loader的__getitem__():

class Mydataset(Dataset):
    def __init__(self, data):
        self.data = data
    
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        return self.data[idx].unsqueeze(-1) # return dimension [Batch, Sequence, Features]

改成這樣得到的pad packed sequence會變成:

tensor([[[1],
         [1],
         [1],
         [1],
         [1],
         [1]],

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

然後放個LSTM model，記得也要設置batch_first。並將我們的tensor轉成float

# LSTM model
net = nn.LSTM(1, 4, 2, batch_first=True) 
out, (h, c) = net(packed_batch_x.float())

pprint(out)
# PackedSequence(
#   data=tensor([[-0.0417, -0.1202,  0.0748,  0.0658],
#         [-0.0434, -0.1283,  0.0659,  0.0617],
#         [-0.0873, -0.1682,  0.1064,  0.0991],
#         [-0.0899, -0.1807,  0.0927,  0.0926],
#         [-0.1233, -0.1861,  0.1222,  0.1136],
#         [-0.1483, -0.1918,  0.1307,  0.1192],
#         [-0.1647, -0.1926,  0.1354,  0.1211],
#         [-0.1752, -0.1915,  0.1380,  0.1214]], grad_fn=<CatBackward>), 
#   batch_sizes=tensor([2, 2, 1, 1, 1, 1]), sorted_indices=None, unsorted_indices=None)

可以看到output也是一個PackedSequence

batch_sizes和packed_batch_x的相同(這也很合理，想想看就知道了)
out的shape是(8, 4)
- packed_batch_x總共有8個非0的data
- LSTM的hidden size為4

如果把out透過pad_packed_sequence()還原的話

pad_out, out_len = rnn_utils.pad_packed_sequence(out, batch_first=True)
pprint(pad_out)
# tensor([[[0.0872, 0.0783, 0.1465, 0.1197],
#          [0.1331, 0.1032, 0.2627, 0.1835],
#          [0.1564, 0.1090, 0.3387, 0.2184],
#          [0.1685, 0.1088, 0.3853, 0.2377],
#          [0.1750, 0.1073, 0.4137, 0.2484],
#          [0.1787, 0.1058, 0.4312, 0.2543]],
# 
#         [[0.0753, 0.0784, 0.1402, 0.1334],
#          [0.1139, 0.1052, 0.2527, 0.2110],
#          [0.0000, 0.0000, 0.0000, 0.0000],
#          [0.0000, 0.0000, 0.0000, 0.0000],
#          [0.0000, 0.0000, 0.0000, 0.0000],
#          [0.0000, 0.0000, 0.0000, 0.0000]]], grad_fn=<TransposeBackward0>)

pprint(out_len)
# tensor([6, 2])

pad_out的shape為(2, 6, 4)
- 2是batch size
- 6是sequence長度，因為變回有padding版本所以每個sequence都一樣長
- 4是LSTM的hidden size
out_len的輸出則是原始padding版本下data的shape

好，完畢！

所以我說…那個pack_sequence勒?

如果你去看source code，你會發現pack_sequence()其實就是padding + pack一起用而已

0
1
2

def pack_sequence(sequences, enforce_sorted=True):
    lengths = torch.as_tensor([v.size(0) for v in sequences])
    return pack_padded_sequence(pad_sequence(sequences), lengths, enforce_sorted=enforce_sorted)

pack_sequence我目前看到比較少在使用(如果有看到也歡迎提供，我會再補上)，大部分都是pad_sequence, pack_padded_sequence, pad_packed_sequence的搭配

然後注意source code都沒寫到batch_first，所以default都是False

延伸討論: batch_size=True到底在幹嘛?

前面提到過:

batch_first=True時，return的維度是: [Batch, Sequence, Features]
batch_first=False時，return的維度是: [Sequence, Batch, Features]

阿這兩個到底差在哪裡，都幾?

我們先放上一張Pytorch doc的圖來幫助理解:

通常第一個維度會被稱之循環維度，也就是機器單次是怎麼抓取和處理數據的

當資料格式是[Batch, Sequence, Features]，代表每次都是抓取同一個sequence內不同timestamp的data
當資料格式是[Sequence, Batch, Features]，代表每次都是抓取同不同sequences下，同一個timestamp的data

而對於後者，很直觀地想到這樣的格式是可以平行化計算的，因為RNN model必須先計算完第t個timestamp的數據後才能計算t+1個數據，所以後者這種cross-batch的方式更加適合。

咦既然batch_first=False比較好，那為啥前面都設置成True呢?

因為Pytorch內做了設置，不管設置True或False他內部都會以[Sequence, Batch, Features]的方式來處理數據(读PyTorch源码学习RNN(1))，看到有些資料說，該參數的設置只是提醒有這個trick

所以在撰寫程式的時候還是以我們平常思考的模式來寫就好。

延伸討論: PackedSequence搭配DataParallel的問題

這是之前實作上遇到的問題，當初有記錄下來，這個問題主要是說:

將data放到不同的gpu上跑時，由於使用了PackedSequence，每個batch的長度都是不固定的，在每張gpu上執行pad_packed_sequence()時，會取它當下batch的最大長度來對其他句子進行padding，這時因為每個gpu上data不同導致當下的最大長度都會不同，在gather的時候就會產生維度不匹配的問題。

關於詳細內容可以看這篇文章: [Pytorch]當DataParallel碰上RNN的那些坑

總結

這篇文章提及了下列幾個問題，可以一邊思考下面的問題來幫助自己確認觀念熟悉了沒:

NLP資料在訓練時長度不同應該怎麼處理?
PackedSequence的概念是什麼? 為何這樣可以增加計算效率?
為什麼pack_padded_sequence()要先對sequences排序?
RNN系列模型常常會有一個參數batch_first，這個參數的作用是什麼?

前言