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

...が, 公式ドキュメントを読んでいるだけではイメージが掴みにくかったため実験してみました。

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

1. pad_sequence

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

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

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などに入力するとエラーとなってしまいます:

# 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にベクトルのリストを与えると, (最大シーケンス長×バッチサイズ×特徴量次元数)のテンソルを適当にゼロ埋めして返してくれます:

print(rnn.pad_sequence([a,b,c]))

# 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にしましょう:

# batch_first=Trueにする
print(rnn.pad_sequence([a,b,c], batch_first=True))

# 第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.]]])

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

# padding_value=-1にする
print(rnn.pad_sequence([a,b,c], batch_first=True, padding_value=-1))

# -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.]]])

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

# 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オブジェクトに格納するための関数です。

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)

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

# 系列長が 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とすれば自動でソートしておいてくれます:

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オブジェクトの形で保管しておいた場合に, サンプルをテンソルに"解凍"するようなイメージかもしれません。

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にしましょう。ゼロ以外で埋めることもできます:

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を指定するといくらでも長い系列に揃えることができる点です:

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オブジェクトへと変換するための関数です。

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.]]])

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

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で系列モデルを訓練する際にサンプルの系列長を揃える関数を見てみました。
ニューラルネットワークの訓練を回す直前のエラーに苦しめられている方の役に少しでも立てば幸いです。