はじめに

こんにちは、とあるIT企業の人間です。
自然言語処理の勉強のためにPytorchを使用して、LSTMモデルを実装していたところ、「お、系列長バラバラやけど、どうすんの？」と思い、自分なりに上手く書けたのでこの記事を書きました。
あくまでも我流の一手法のため、参考にしていただけるとありがたいです。

なお、関数などの詳細は公式ページを参考にしてください。

何かご指摘・質問等ございましたら、ご遠慮なくコメント欄にお願い致します。

目的

自然言語をIDで表現した状態のデータがある場合に以下の処理をスマートに実装したい

系列長がバラバラのため、パディングする。
Embedding層、LSTM層に入力する（順伝搬計算）
それぞれの系列長の最終出力を抽出する。

つまり以下のような系列長が違うデータをスマートな実装で一気に処理したい、ということです。
データ例：
inp = [[ 7, 3, 5], [ 9, 5], [ 7, 1, 8, 4, 4]]

環境

Python 3.7.6
pytorch 1.5.0

実装

##データ
系列長がバラバラのリストと、それぞれの系列長が入ったリストを準備します。
リストで各要素をtensor型に変換してください。今回では以下のデータを準備します。

.py

# 使用データの定義
word_id = [T.tensor([ 7, 3, 5]),        # 3
           T.tensor([ 9, 5]),           # 2
           T.tensor([ 7, 1, 8, 4, 4])]  # 5

こちらの系列長は、$[3,2,5]$です。こちらもリストで用意します。

.py

# 使用データの定義
seq_list = [3, 2, 5]

##1. 系列長がバラバラのため、パディングする。
pytorchには、パディングしてくれる関数pad_packed_sequenceがあるので、そちらを使いますが、その入力のためにまずpack_sequenceを使用します。

.py

# パディング
## ①PackedSequenceという一つのオブジェクト型に変換する。
word_id_packed = nn.utils.rnn.pack_sequence(word_id, enforce_sorted=False)
## ②0埋めする
word_id_padded, _ = nn.utils.rnn.pad_packed_sequence(word_id_packed,batch_first =True,padding_value=0)

word_id_paddedの出力結果です。
お、0埋めできてますね。

tensor([[7, 3, 5, 0, 0],
        [9, 5, 0, 0, 0],
        [7, 1, 8, 4, 4]])

「0以外でパディングしたい」「最大系列長を指定したい」と思う方は公式ページに引数の説明があるのでそちらを見ていただくと良いかと思います。

##2. Embedding層、LSTM層に入力する（順伝搬計算）
ここでは最終的に各系列長の出力が欲しいので、パディングした状態を保ったままLSTM層に入力する必要があります。そのため、LSTM層への入力前にpack_padded_sequenceという関数を使用します。

また「"0"はパディングしたもの」として扱いたいので、nn.Embeddingでpadding_idx=0を指定します。"-1"など他の値でパディングした場合は、padding_idxの値を変更してください。

.py

# 各層のパラメータ定義
## Embedding層（辞書の単語数：10、次元数：5）
emb_layer = nn.Embedding(10, 5, padding_idx=0)
## LSTM層（入力次元数：5、出力次元数：2）
lstm_layer = nn.LSTM(5, 2) 


# 順伝搬計算
## Embedding層へ入力
word_emb = emb_layer(word_id_padded)
## パディングした状態を保持したまま、PackedSequenceオブジェクト型に変換する。
word_emb_packed = nn.utils.rnn.pack_padded_sequence(word_emb, seq_list, batch_first =True, enforce_sorted=False)
## LSTM層へ入力する。この時、第二返り値も受け取るようにセットする。(ht, ct)
lstm, (ht, ct)  = lstm_layer(word_emb_packed)

lstmの出力結果です。
各系列長分の計算だけができていますね。Good。

tensor([[[ 0.2595,  0.0771],
         [ 0.2072,  0.2954],
         [ 0.3157,  0.0541],
         [ 0.0000,  0.0000],
         [ 0.0000,  0.0000]],

        [[ 0.3509,  0.2226],
         [ 0.2906,  0.0636],
         [ 0.0000,  0.0000],
         [ 0.0000,  0.0000],
         [ 0.0000,  0.0000]],

        [[ 0.2595,  0.0771],
         [ 0.2382,  0.0205],
         [ 0.2692,  0.1638],
         [ 0.2793,  0.0356],
         [ 0.3055, -0.0184]]]
          , grad_fn=<IndexSelectBackward>)

##3. それぞれの系列長の最終出力を抽出する
最後に先ほどのhtを見てみましょう。

.py

print(ht[-1])

出力結果です。これだけ見ても「お？」という感じなので先ほどの出力を見てみましょう。

tensor([[ 0.3157,  0.0541],
        [ 0.2906,  0.0636],
        [ 0.3055, -0.0184]], grad_fn=<SelectBackward>)

おお！！lstmの出力からうまく抽出できてますね！！

tensor([[[ 0.2595,  0.0771],
         [ 0.2072,  0.2954],
         [ 0.3157,  0.0541],  <------
         [ 0.0000,  0.0000],
         [ 0.0000,  0.0000]],

        [[ 0.3509,  0.2226],
         [ 0.2906,  0.0636],  <------
         [ 0.0000,  0.0000],
         [ 0.0000,  0.0000],
         [ 0.0000,  0.0000]],

        [[ 0.2595,  0.0771],
         [ 0.2382,  0.0205],
         [ 0.2692,  0.1638],
         [ 0.2793,  0.0356],
         [ 0.3055, -0.0184]]] <------
          , grad_fn=<IndexSelectBackward>)

全体のコード

個人的にはすっきりしました。

padding_lstm.py

# 使用データの定義
word_id = [T.tensor([ 7, 3, 5]),        # 3
           T.tensor([ 9, 5]),           # 2
           T.tensor([ 7, 1, 8, 4, 4])]  # 5

seq_list = [3, 2, 5]


# パディング
## ①PackedSequenceという一つのオブジェクト型に変換する。
word_id_packed = nn.utils.rnn.pack_sequence(word_id, enforce_sorted=False)
## ②0埋めする
word_id_padded, _ = nn.utils.rnn.pad_packed_sequence(word_id_packed,batch_first =True,padding_value=0)


# 各層のパラメータ定義
## Embedding層（辞書の単語数：10、次元数：5）
emb_layer = nn.Embedding(10, 5, padding_idx=0)
## LSTM層（入力次元数：5、出力次元数：2）
lstm_layer = nn.LSTM(5, 2) 


# 順伝搬計算
## Embedding層へ入力
word_emb = emb_layer(word_id_padded)
## パディングした状態を保持したまま、PackedSequenceオブジェクト型に変換する。
word_emb_packed = nn.utils.rnn.pack_padded_sequence(word_emb, seq_list, batch_first =True, enforce_sorted=False)
## LSTM層へ入力する。この時、第二返り値も受け取るようにセットする。(ht, ct)
lstm, (ht, ct)  = lstm_layer(word_emb_packed)

# 
print(ht[-1])