💨

PackedSequenceオブジェクトを理解してpadding処理をマスターする

2023/02/19に公開

この記事の概要

GraphGen(GitHub)のコードを読んでいて、paddingの処理が理解できず詰まりました。自分の知りたいことをピンポイントで紹介している記事がなかったため、今後同じようなことで詰まった人のために記事にしようと考えました。
最初に結論を述べると、PyTorchにおけるpaddingでわかりづらいところは、PackedSequenceオブジェクトです。これは、元データとそれぞれのデータ長をセットで管理するデータ構造です。可変長時系列データの処理では、自分の持っているデータをpadding処理する必要があります。しかし、データの内容や形状によっては用意されている関数でpaddingできないこともあります。そこでPyTorchでは、一旦データをPackedSequenceオブジェクトに変換する処理(packing処理)を間に挟むことで柔軟にpadding処理を行えるようになっています。
一度理解してしまえばPyTorchの細やかな設計に感服するばかりですが、最初は何が何だかわからずむしろ混乱します。そこで、本記事ではPyTorchでpadding処理のために用意されている4つの関数を紹介し、その使い方の例を示したいと思います。

PyTorchに用意されているpaddingに関する4つの関数

以下の関数は、torch.nn.utils.rnnからインポートできます。

  • pack_sequence(公式ドキュメント)
    • 長さの揃っていないtensorをリスト形式で持つデータをPackedSequenceオブジェクトに変換します。
  • pack_padded_sequence(公式ドキュメント)
    • 既に何らかの値でpaddingされ、長さは揃っているtensorをPackedSequenceオブジェクトに変換します。
  • pad_sequence(公式ドキュメント)
    • 長さの揃っていないtensorをリスト形式で持つデータを長さが揃うように0でpaddingを行います。
  • pad_packed_sequence(公式ドキュメント)
    • PackedSequenceオブジェクトを受け取って、元データを0でpaddingしたものを返します。

PackedSequenceオブジェクトについて

PackedSequenceオブジェクトとは、PyTorchで用意されているデータ構造の一つです。このオブジェクトは、すべてのRNNモジュールのinputとして用いることができます。このオブジェクトが保持するのは、入力データとそれらのバッチサイズのリストです。これら二つデータからPyTorchは自動的に元のデータを0でpaddingしたものを生成することができます。
以下にPackedSequenceオブジェクトの例を示します。

PackedSequence(data=tensor([4, 2, 1, 5, 3, 6]), 
              batch_sizes=tensor([3, 2, 1]), 
              sorted_indices=None, 
              unsorted_indices=None)
  • data: 元データを1次元のtensorで保持しています。
  • batch_sizes: 1次元で保持しているtensorを元データに復元するためにそれぞれのデータ長を保持しています。
  • sorted_indices: 元データからどのようにしてPackedSequenceオブジェクトが生成されたのかを保持するtensorです。入力データがソート済みでない場合、入力されたデータは自動的にソートされます。ソート済みデータのindexとPackedSequenceオブジェクトの関係を保持します。
  • unsorted_indices: こちらも上とほぼ同様ですが、PackedSequenceオブジェクトからソート済みの元データを復元した後に、元のデータ順に戻すためのindexを保持しています。

pad_sequenceの使い所

小規模なコードを書くときは、この関数でpaddingすることが多いのではないかと思います。取得したデータをそれぞれtensorに変換したものをリスト形式にするだけで前処理が終わります。あとはpadding_valueで指定した値でpadding処理ができます。サンプルコードでは、見やすさのためにbatch_first=Trueとしています。これ以降のサンプルコードも同様です。

# pad_sequenceの使い所
# 長さの異なるtensorをリスト形式で持っている時に用いると便利

from torch.nn.utils.rnn import pad_sequence
import torch

a = torch.tensor([4,5,6])
b = torch.tensor([2,3])
c = torch.tensor([1])
input_data = ([a,b,c])
print(input_data)
# [tensor([1]), tensor([2, 3]), tensor([4, 5, 6])]
padded = pad_sequence(input_data, batch_first=True, padding_value=0)
print(padded)
# tensor([[4, 5, 6],
#         [2, 3, 0],
#         [1, 0, 0]])

pack_sequenceの使い所

用いる入力データは、pad_sequenceの時に使ったinput_dataと同様です。pack_sequenceinput_dataenforced_sorted:boolを与えると入力データがPackedSequenceオブジェクトに変換されます。この操作のことをpackingすると呼ぶことにします。
enforced_sorted:boolとは、input_dateのリストに含まれるtensorの長さが降順になっていればTrue、なっていないならばFalseを与えます。デフォルトがTrueなのでサンプルコードでは、tensorの長さが降順になるようにinput_dataを定義しました。

PackedSequenceの復元

pack_sequencePackedSequenceオブジェクトを生成したあと、pad_packed_sequenceで元のデータをpaddingしたものを取得します。paddingの値は、padding_valueで設定できます。また、total_lengthは出力するデータの長さを定める値です。今回用いたデータの最大長は3なので3以上の任意の値を選択できます。選択した値に従って、その長さまでpaddingを行います。

pack_sequenceを用いるメリット

一度PackedSequenceオブジェクトを経由してからpad_packed_sequenceを用いてデータをpaddingすることで任意の長さの出力を得ることができます。予めモデル側の入力長が定められている場合、total_lengthを変更するだけで入力データの調整ができます。とはいえ、自分はそれが必要になったことがないので実例があれば知りたいです。

# pack_sequenceの使い所
# 長さの異なるtensorをリスト形式で得ることができ、
# RNNの入力としてPackedSequenceオブジェクトを用いたいときに使えます。
# また、出力の長さを任意に設定したい時にも有効です。
# pad_sequenceでは、自動的に入力データの最大長に合わせて
# paddingが行われるため任意の長さのデータを生成することができません。
from torch.nn.utils.rnn import pack_sequence, pad_packed_sequence
import torch

a = torch.tensor([4,5,6])
b = torch.tensor([2,3])
c = torch.tensor([1])
input_data = ([a,b,c])
print(input_data)
# [tensor([1]), tensor([2, 3]), tensor([4, 5, 6])]
packed = pack_sequence(input_data, enforce_sorted=True) # enforce_sortedはデフォルトがTrue
print(packed)
# PackedSequence(data=tensor([4, 2, 1, 5, 3, 6]), 
# batch_sizes=tensor([3, 2, 1]), 
# sorted_indices=None, unsorted_indices=None)
padded, _ = pad_packed_sequence(packed, batch_first=True, 
                                padding_value=0, total_length=None)
print(padded)
# tensor([[4, 5, 6],
#         [2, 3, 0],
#         [1, 0, 0]])

pack_padded_sequenceの使い所

既に何らかの値でpaddingされ長さが揃えられた上で一つのtensorにtorch.catされているような入力データをpackingする際に有効です。pack_padded_sequenceinput_data及びデータ長を格納したtensor input_lenを与えることで、PackedSequenceオブジェクトを生成してくれます。この時、重要なのはデータ長を保持するtensorを予め作成する必要があり、それが実際にinput_dataに入っているデータの順序に一致している必要があるということです。ここを間違えると結果に何らかの影響が出ますが、デバック時に気づくのは難しいと思いますので、実装時点で細心の注意を払いましょう。

PackedSequenceの復元

pack_padded_sequenceが生成するのは、pack_sequenceと同じくPackedSequenceオブジェクトなので復元手順は全く同一です。これが統一したデータ構造を利用するメリットですね。使いやすいライブラリを提供してくれる開発者には感謝しかありません。

# pack_padded_sequenceの使い所
# 既に何らかの値でpaddingされた複数のデータを持つtensorを
# paddingしたい時はこれを用います。
# GraphGenがこれに該当しました。
from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence
import torch

a = torch.tensor([4,5,6,-1]) # paddingの値を-1とします
b = torch.tensor([2,3,-2,-2]) # paddingの値を-2とします
c = torch.tensor([1,-3,-3,-3]) # paddingの値を-3とします
input_len = torch.tensor([3,2,1]) # a,b,cそれぞれのデータ長を保持するリスト
input_data = torch.vstack((a,b,c)) # tensorを結合
print(input_data)
# tensor([[ 4,  5,  6, -1],
#         [ 2,  3, -2, -2],
#         [ 1, -3, -3, -3]])
packed = pack_padded_sequence(input_data, input_len,
                             batch_first=True, enforce_sorted=True)
print(packed)
# PackedSequence(data=tensor([4, 2, 1, 5, 3, 6]), 
# batch_sizes=tensor([3, 2, 1]), sorted_indices=None, 
# unsorted_indices=None)
padded, _ = pad_packed_sequence(packed, batch_first=True, 
                                padding_value=0, total_length=None)
print(padded)
# tensor([[4, 5, 6],
#         [2, 3, 0],
#         [1, 0, 0]])

最後に

本記事では、PyTorchので用いられるPackedSequenceオブジェクトについて解説しました。また、その過程でPyTorchにおけるpaddingやpackingについてサンプルコードを示しながら説明を行いました。この記事だけで、PyTorchのpacking処理及びpadding処理についてまとまった知識を得られるように工夫したつもりです。参考になった部分や欠けている部分があれば教えていただけると幸いです。

参考記事

GitHubで編集を提案

Discussion