PyTorchのシーケンスのpaddingとpackingを使い分ける

## 0. はじめに

PyTorchでRNN, LSTM, GRUなどの系列モデルを訓練するには, サンプルの系列としての長さが全て同じでなければなりません。

(ニューラルネットワークの仕組み的にはそんな必要はありませんが, 実際のコーディングでは系列長がバラバラだとエラーを吐いてしまいます)

しかし, 訓練させたい文章サンプルが最初から全部同じ長さであることはまず無いでしょう。

このため, 

- 短い文章の末尾にダミー単語を加えておく
- 長い文章の尻尾を切り取っておく

などの工夫をして, **事前にサンプルの系列長を揃える処理をしておく**必要があります。

ありがたいことに, PyTorchは系列長を揃えてくれるような関数を4つ用意してくれています:

- `torch.nn.utils.rnn.pad_sequence`
- `torch.nn.utils.rnn.pack_sequence`
- `torch.nn.utils.rnn.pad_packed_sequence`
- `torch.nn.utils.rnn.pack_padded_sequence`

...が, [公式ドキュメント](https://pytorch.org/docs/stable/nn.html#torch.nn.utils.rnn.pack_sequence)を読んでいるだけではイメージが掴みにくかったため実験してみました。

ちなみに, 今のうちに言ってしまいますが**1個目の`torch.nn.utils.rnn.pad_sequence`だけ覚えてしまえば実際の機械学習で困ることはほぼないです**。

## 1. pad_sequence

長さの違うテンソルを与えると, 短いものの末尾にゼロ埋めを施して次元を揃えてくれる関数です。

次のようなベクトルa,b,cをRNNなどに入力したいとしましょう:

```Python3
import torch
import torch.nn.utils.rnn as rnn

# a, b, cの系列長はバラバラ
# 系列長3, 特徴量5次元
a = torch.ones(3, 5)
# 系列長2, 特徴量5次元
b = torch.ones(2, 5) * 2
# 系列長1, 特徴量5次元
c = torch.ones(1, 5) * 3
```

系列長が揃っていないので, 1回の訓練でこれら3つをまとめてRNNなどに入力するとエラーとなってしまいます:

```Python3
# RNNモデルをつくる (隠れ状態を仮に6次元とする)
# 注意: RNNモデル自体をつくるメソッドでは'RNN'を大文字表記する
mymodel = torch.nn.RNN(input_size=5, hidden_size=6, batch_first=True)

# a, b, cを別々に入力するのはOK
batch_size = 1
h0 = torch.randn(1, batch_size, 6)
h_a, hn_a = mymodel(a, h0)
h_b, hn_b = mymodel(b, h0)
h_c, hn_c = mymodel(c, h0)

# a, b, cをまとめて入力するとエラー
batch_size = 3
h0 = torch.randn(1, batch_size, 6)
h, hn = mymodel(torch.tensor([a,b,c]), h0)  # ERROR
```

`torch.nn.utils.rnn.pad_sequence`にベクトルのリストを与えると, (最大シーケンス長×バッチサイズ×特徴量次元数)のテンソルを適当にゼロ埋めして返してくれます:

```Python3
print(rnn.pad_sequence([a,b,c]))
```
```Python3
# a, b, cの系列長がすべて3になった
tensor([[[1., 1., 1., 1., 1.],
         [2., 2., 2., 2., 2.],
         [3., 3., 3., 3., 3.]],

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

        [[1., 1., 1., 1., 1.],
         [0., 0., 0., 0., 0.],
         [0., 0., 0., 0., 0.]]])
```


軸の並びを(バッチサイズ×最大シーケンス長×特徴量次元数)にしたければ`batch_first=True`にしましょう:

```Python3
# batch_first=Trueにする
print(rnn.pad_sequence([a,b,c], batch_first=True))
```
```Python3
# 第1軸がサンプル方向になった
tensor([[[1., 1., 1., 1., 1.],
         [1., 1., 1., 1., 1.],
         [1., 1., 1., 1., 1.]],

        [[2., 2., 2., 2., 2.],
         [2., 2., 2., 2., 2.],
         [0., 0., 0., 0., 0.]],

        [[3., 3., 3., 3., 3.],
         [0., 0., 0., 0., 0.],
         [0., 0., 0., 0., 0.]]])
```

ゼロ以外の値で埋めることもできます:

```Python3
# padding_value=-1にする
print(rnn.pad_sequence([a,b,c], batch_first=True, padding_value=-1))
```
```Python3
# -1でパディングされた
tensor([[[ 1.,  1.,  1.,  1.,  1.],
         [ 1.,  1.,  1.,  1.,  1.],
         [ 1.,  1.,  1.,  1.,  1.]],

        [[ 2.,  2.,  2.,  2.,  2.],
         [ 2.,  2.,  2.,  2.,  2.],
         [-1., -1., -1., -1., -1.]],

        [[ 3.,  3.,  3.,  3.,  3.],
         [-1., -1., -1., -1., -1.],
         [-1., -1., -1., -1., -1.]]])
```

これで無事にサンプルの系列長を揃えることができました。

```Python3
# a, b, cをまとめてRNNに入力できるようになった!
batch_size = 3
h0 = torch.randn(1, batch_size, 6)
c0 = torch.randn(1, batch_size, 6)
X = rnn.pad_sequence([a,b,c], batch_first=True)
h, (hn, cn) = mymodel(X, (h0, c0))
```


## 2. pack_sequence

**ここから先は読み飛ばしても構いません**(実際の機械学習で使うことがほとんどありません)。

なんだか上の`pad_sequence`さえ使えば万事解決してしまいそうな予感がしますが, どんな機能なのか確認してみましょう。

複数のサンプルを1つにまとめて`PackedSequence`オブジェクトに格納するための関数です。


```Python3
print(rnn.pack_sequence([a,b,c]))
```

```
PackedSequence(data=tensor([[1., 1., 1., 1., 1.],
        [2., 2., 2., 2., 2.],
        [3., 3., 3., 3., 3.],
        [1., 1., 1., 1., 1.],
        [2., 2., 2., 2., 2.],
        [1., 1., 1., 1., 1.]]), batch_sizes=tensor([3, 2, 1]), sorted_indices=None, unsorted_indices=None)
```

このとき, 引数のリストがサンプルの系列長に関して降順ソートされていないとエラーとなります:

```Python3
# 系列長が b>=c>=a ではないのでエラー
print(rnn.pack_sequence([b,c,a]))
```
```
---------------------------------------------------------------------------
RuntimeError                              

    (中略)

RuntimeError: `lengths` array must be sorted in decreasing order when `enforce_sorted` is True. You can pass `enforce_sorted=False` to pack_padded_sequence and/or pack_sequence to sidestep this requirement if you do not need ONNX exportability.
```

しかし`enforce_sorted=False`とすれば自動でソートしておいてくれます:

```Python3
print(rnn.pack_sequence([b,c,a], enforce_sorted=False))
```

```
PackedSequence(data=tensor([[1., 1., 1., 1., 1.],
        [2., 2., 2., 2., 2.],
        [3., 3., 3., 3., 3.],
        [1., 1., 1., 1., 1.],
        [2., 2., 2., 2., 2.],
        [1., 1., 1., 1., 1.]]), batch_sizes=tensor([3, 2, 1]), sorted_indices=tensor([2, 0, 1]), unsorted_indices=tensor([1, 2, 0]))
```

ところで, `PackedSequence`オブジェクトは`torch.nn.utils.rnn.pack_sequence`関数の戻り値としてのみ取得することができ, コンストラクタで生成することはできないようです。



## 3. pad_packed_sequence
先ほどの`PackedSequence`オブジェクトを入れると, ゼロ埋めして系列長の揃ったテンソルを返してくれる関数です。

あらかじめサンプル情報を素のベクトルではなく`PackedSequence`オブジェクトの形で保管しておいた場合に, サンプルをテンソルに"解凍"するようなイメージかもしれません。



```Python3
packed = rnn.pack_sequence([a,b,c])

batch, len_batch = rnn.pad_packed_sequence(packed)

print(batch)
print(len_batch)
```

戻り値はタプルで, (最大シーケンス長×バッチサイズ×特徴量次元数)のテンソルと, 各サンプルの系列長を容れたベクトルが与えられます:

```
tensor([[[1., 1., 1., 1., 1.],
         [2., 2., 2., 2., 2.],
         [3., 3., 3., 3., 3.]],

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

        [[1., 1., 1., 1., 1.],
         [0., 0., 0., 0., 0.],
         [0., 0., 0., 0., 0.]]])
```

```
tensor([3, 2, 1])
```

軸の並びを(バッチサイズ×最大シーケンス長×特徴量次元数)にしたければ`batch_first=True`にしましょう。ゼロ以外で埋めることもできます:

```Python3
batch_2, len_batch_2 = rnn.pad_packed_sequence(packed, batch_first=True, padding_value=-1)

print(batch_2)
print(len_batch_2)
```

```
tensor([[[ 1.,  1.,  1.,  1.,  1.],
         [ 1.,  1.,  1.,  1.,  1.],
         [ 1.,  1.,  1.,  1.,  1.]],

        [[ 2.,  2.,  2.,  2.,  2.],
         [ 2.,  2.,  2.,  2.,  2.],
         [-1., -1., -1., -1., -1.]],

        [[ 3.,  3.,  3.,  3.,  3.],
         [-1., -1., -1., -1., -1.],
         [-1., -1., -1., -1., -1.]]])
```
```
tensor([3, 2, 1])
```

なんだか最初にみた`pad_sequence`と同じことをわざわざ遠回しにやっているだけに見えますね?

ただの`pad_sequence`との最大の違いは, `total_length`を指定するといくらでも長い系列に揃えることができる点です:

```Python3
batch_3, len_batch_3 = rnn.pad_packed_sequence(packed, batch_first=True, total_length=6)

print(batch_3)
```

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

        [[2., 2., 2., 2., 2.],
         [2., 2., 2., 2., 2.],
         [0., 0., 0., 0., 0.],
         [0., 0., 0., 0., 0.],
         [0., 0., 0., 0., 0.],
         [0., 0., 0., 0., 0.]],

        [[3., 3., 3., 3., 3.],
         [0., 0., 0., 0., 0.],
         [0., 0., 0., 0., 0.],
         [0., 0., 0., 0., 0.],
         [0., 0., 0., 0., 0.],
         [0., 0., 0., 0., 0.]]])
```

## 4. pack_padded_sequence

`pad_sequence`した結果を`PackedSequence`オブジェクトへと変換するための関数です。

```Python3
padded = rnn.pad_sequence([a,b,c], batch_first=True)
print(padded)
```
```
tensor([[[1., 1., 1., 1., 1.],
         [1., 1., 1., 1., 1.],
         [1., 1., 1., 1., 1.]],

        [[2., 2., 2., 2., 2.],
         [2., 2., 2., 2., 2.],
         [0., 0., 0., 0., 0.]],

        [[3., 3., 3., 3., 3.],
         [0., 0., 0., 0., 0.],
         [0., 0., 0., 0., 0.]]])
```

引数にはパディング済みのテンソルのほか, 各サンプルの系列長をリストで与える必要があります:

```Python3
print(rnn.pack_padded_sequence(padded, [3, 2, 1]))
```
```
PackedSequence(data=tensor([[1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [2., 2., 2., 2., 2.],
        [2., 2., 2., 2., 2.],
        [3., 3., 3., 3., 3.]]), batch_sizes=tensor([3, 2, 1]), sorted_indices=None, unsorted_indices=None)
```

するとパディングしてあったサンプルがパディング前の系列長に"切り取られて"保管されていることがわかります。
今回扱った4つの関数のうち一番使いどころが少ないかもしれません...。


## 5. おわりに
以上, PyTorchで系列モデルを訓練する際にサンプルの系列長を揃える関数を見てみました。
ニューラルネットワークの訓練を回す直前のエラーに苦しめられている方の役に少しでも立てば幸いです。