Keras でマスキングとパディングをする

TensorFlow.orgで表示 Google Colab で実行 GitHub でソースを表示 ノートブックをダウンロード

MNIST モデルをビルドする

import numpy as np
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
2024-01-11 19:57:50.441698: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
2024-01-11 19:57:50.441746: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
2024-01-11 19:57:50.443298: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered

はじめに

マスキングは、シーケンス処理レイヤーに入力の特定の時間ステップが欠落しているためデータを処理する際にスキップする必要があることを伝えるために使用する手法です。

パディングは、マスキングされたステップがシーケンスの先頭または末尾にある特殊なマスキングです。パディングは、シーケンスデータを連続したバッチにエンコードする必要性から生まれました。バッチ内のすべてのシーケンスを所定の標準の長さに合わせるためには、一部のシーケンスをパディングまたはトランケートする(切り詰める)必要があるためです。

では、詳しく見ていきましょう。

パディングシーケンスデータ

シーケンスデータを処理する際に個々のサンプルの長さが異なることは、非常に一般的です。次の例(単語としてトークン化されたテキスト)を考えてみます。

[
  ["Hello", "world", "!"],
  ["How", "are", "you", "doing", "today"],
  ["The", "weather", "will", "be", "nice", "tomorrow"],
]

語彙検索の後、データは以下のように整数としてベクトル化されるかもしれません。

[
  [71, 1331, 4231]
  [73, 8, 3215, 55, 927],
  [83, 91, 1, 645, 1253, 927],
]

データは、個々のサンプルがそれぞれ 3、5、6 の長さを持つネストされたリストです。ディープラーニングモデルの入力データは,単一のテンソル(例えばこの場合だと(batch_size, 6, vocab_size)のような形状)でなければならないため、最長のアイテムよりも短いサンプルは、何らかのプレースホルダー値でパディングする必要があります。(その代わりに、短いサンプルをパディングする前に長いサンプルをトランケートすることも可能です。)

Keras は Python のリストを共通の長さにトランケートしたりパディングしたりするユーティリティ関数を提供します:tf.keras.preprocessing.sequence.pad_sequences

raw_inputs = [
    [711, 632, 71],
    [73, 8, 3215, 55, 927],
    [83, 91, 1, 645, 1253, 927],
]

# By default, this will pad using 0s; it is configurable via the
# "value" parameter.
# Note that you could "pre" padding (at the beginning) or
# "post" padding (at the end).
# We recommend using "post" padding when working with RNN layers
# (in order to be able to use the
# CuDNN implementation of the layers).
padded_inputs = tf.keras.preprocessing.sequence.pad_sequences(
    raw_inputs, padding="post"
)
print(padded_inputs)
[[ 711  632   71    0    0    0]
 [  73    8 3215   55  927    0]
 [  83   91    1  645 1253  927]]

マスキング

全てのサンプルが統一された長さになったので、今度はデータの一部が実際にパディングされ、無視されるべきであることをモデルに知らせなければなりません。このメカニズムがマスキングです。

Keras モデルで入力マスクを導入するには、3 つの方法があります。

  • keras.layers.Masking レイヤーを追加する。
  • keras.layers.Embedding レイヤーを mask_zero=True で設定する。
  • mask引数をサポートするレイヤー(RNN レイヤーなど)を呼び出す際に、この引数を手動で渡す。

マスク生成レイヤー : EmbeddingMasking

内部でこれらのレイヤーはマスクテンソル(形状(batch, sequence_length)の 2 次元テンソル)を作成し、Masking または Embedding レイヤーによって返されるテンソル出力にアタッチします。

embedding = layers.Embedding(input_dim=5000, output_dim=16, mask_zero=True)
masked_output = embedding(padded_inputs)

print(masked_output._keras_mask)

masking_layer = layers.Masking()
# Simulate the embedding lookup by expanding the 2D input to 3D,
# with embedding dimension of 10.
unmasked_embedding = tf.cast(
    tf.tile(tf.expand_dims(padded_inputs, axis=-1), [1, 1, 10]), tf.float32
)

masked_embedding = masking_layer(unmasked_embedding)
print(masked_embedding._keras_mask)
tf.Tensor(
[[ True  True  True False False False]
 [ True  True  True  True  True False]
 [ True  True  True  True  True  True]], shape=(3, 6), dtype=bool)
tf.Tensor(
[[ True  True  True False False False]
 [ True  True  True  True  True False]
 [ True  True  True  True  True  True]], shape=(3, 6), dtype=bool)

出力された結果から分かるように、マスクは形状が(batch_size, sequence_length)の 2 次元ブールテンソルであり、そこでは個々の False エントリは、対応する時間ステップを処理中に無視すべきであることを示しています。

Functional API と Sequential API のマスク伝播

Functional API または Sequential API を使用する場合、Embedding レイヤーまたは Masking レイヤーによって生成されたマスクは、それらを使用できる任意のレイヤー(例えば RNN レイヤーなど)にネットワークを介して伝播されます。Keras は入力に対応するマスクを自動的に取得し、その使用方法を知っている任意のレイヤーに渡します。

例えば、以下の Sequential API モデルでは、LSTM レイヤーは自動的にマスクを取得します。つまりこれは、パディングされた値を無視するということです。

model = keras.Sequential(
    [layers.Embedding(input_dim=5000, output_dim=16, mask_zero=True), layers.LSTM(32),]
)

これは、以下の Functional API モデルでも同様です。

inputs = keras.Input(shape=(None,), dtype="int32")
x = layers.Embedding(input_dim=5000, output_dim=16, mask_zero=True)(inputs)
outputs = layers.LSTM(32)(x)

model = keras.Model(inputs, outputs)

マスクテンソルを直接レイヤーに渡す

マスクを扱うことができるレイヤー(LSTM レイヤーなど)は、それらの __call__ メソッドに mask 引数を持っています。

一方、マスクを生成するレイヤー(例えば
Embedding)は、呼び出し可能な compute_mask(input, previous_mask) メソッドを公開します。

例えば下記のようにして、マスクを生成するレイヤーの compute_mask() メソッドの出力を、マスクを消費するレイヤーの __call__ メソッドに渡すことができます。

class MyLayer(layers.Layer):
    def __init__(self, **kwargs):
        super(MyLayer, self).__init__(**kwargs)
        self.embedding = layers.Embedding(input_dim=5000, output_dim=16, mask_zero=True)
        self.lstm = layers.LSTM(32)

    def call(self, inputs):
        x = self.embedding(inputs)
        # Note that you could also prepare a `mask` tensor manually.
        # It only needs to be a boolean tensor
        # with the right shape, i.e. (batch_size, timesteps).
        mask = self.embedding.compute_mask(inputs)
        output = self.lstm(x, mask=mask)  # The layer will ignore the masked values
        return output


layer = MyLayer()
x = np.random.random((32, 10)) * 100
x = x.astype("int32")
layer(x)
<tf.Tensor: shape=(32, 32), dtype=float32, numpy=
array([[-0.00217493,  0.00870568, -0.00231357, ...,  0.00161815,
         0.00606224,  0.00128512],
       [ 0.00338142,  0.00302714,  0.00811519, ...,  0.00560204,
        -0.00048769, -0.00109794],
       [-0.00092897, -0.00211617, -0.01060253, ...,  0.00087894,
        -0.00760199, -0.00264295],
       ...,
       [-0.00209734,  0.01447711,  0.00503664, ...,  0.0167399 ,
         0.00576398, -0.00300557],
       [ 0.00198119,  0.00383252,  0.00542844, ...,  0.00765714,
         0.0025847 , -0.00153534],
       [-0.00557906, -0.00217974,  0.00527243, ...,  0.00245354,
         0.00294105,  0.00245563]], dtype=float32)>

カスタムレイヤーでマスキングをサポートする

場合によっては、マスクを生成するレイヤー(Embedding など)や、現在のマスクを変更するレイヤーを書く必要があります。

例えば、時間次元で連結する Concatenate レイヤーのように、入力とは異なる時間次元を持つテンソルを生成するレイヤーは、現在のマスクを変更して、マスクされた時間ステップを下流のレイヤーが適切に考慮に入れられるようにする必要があります。

これを行うには、レイヤーに layer.compute_mask() メソッドを実装します。これは、入力と現在のマスクが与えられた時に新しいマスクを生成します。

ここでは、現在のマスクを変更する必要がある TemporalSplit レイヤーの例を示します。

class TemporalSplit(keras.layers.Layer):
    """Split the input tensor into 2 tensors along the time dimension."""

    def call(self, inputs):
        # Expect the input to be 3D and mask to be 2D, split the input tensor into 2
        # subtensors along the time axis (axis 1).
        return tf.split(inputs, 2, axis=1)

    def compute_mask(self, inputs, mask=None):
        # Also split the mask into 2 if it presents.
        if mask is None:
            return None
        return tf.split(mask, 2, axis=1)


first_half, second_half = TemporalSplit()(masked_embedding)
print(first_half._keras_mask)
print(second_half._keras_mask)
tf.Tensor(
[[ True  True  True]
 [ True  True  True]
 [ True  True  True]], shape=(3, 3), dtype=bool)
tf.Tensor(
[[False False False]
 [ True  True False]
 [ True  True  True]], shape=(3, 3), dtype=bool)

もう 1 つの例として、入力値からマスクを生成できる CustomEmbedding レイヤーの例を示します。

class CustomEmbedding(keras.layers.Layer):
    def __init__(self, input_dim, output_dim, mask_zero=False, **kwargs):
        super(CustomEmbedding, self).__init__(**kwargs)
        self.input_dim = input_dim
        self.output_dim = output_dim
        self.mask_zero = mask_zero

    def build(self, input_shape):
        self.embeddings = self.add_weight(
            shape=(self.input_dim, self.output_dim),
            initializer="random_normal",
            dtype="float32",
        )

    def call(self, inputs):
        return tf.nn.embedding_lookup(self.embeddings, inputs)

    def compute_mask(self, inputs, mask=None):
        if not self.mask_zero:
            return None
        return tf.not_equal(inputs, 0)


layer = CustomEmbedding(10, 32, mask_zero=True)
x = np.random.random((3, 10)) * 9
x = x.astype("int32")

y = layer(x)
mask = layer.compute_mask(x)

print(mask)
tf.Tensor(
[[False  True  True  True  True  True  True  True  True  True]
 [ True False  True  True  True  True  True False  True  True]
 [ True  True  True False  True  True  True False  True  True]], shape=(3, 10), dtype=bool)

オプトインして互換性のあるレイヤー間でマスクを伝播する

ほとんどのレイヤーは時間次元を変更しないため、現在のマスクを変更する必要はありません。しかし、現在のマスクを変更せずにそれらを次のレイヤーに伝播したい場合があります。これはオプトイン動作です。 デフォルトでは、(フレームワークがマスクの伝播が安全かどうか判断する方法を持たないため)カスタムレイヤーは現在のマスクを破棄します。

時間次元を変更しないカスタムレイヤーを持ち、それが現在の入力マスクを伝播できるようにしたい場合は、レイヤーのコンストラクタを self.supports_masking = True に設定する必要があります。この場合、compute_mask() のデフォルトの動作は、現在のマスクを通過させるだけとなります。

マスク伝搬のためにホワイトリスト化されたレイヤーの例を示します。:

class MyActivation(keras.layers.Layer):
    def __init__(self, **kwargs):
        super(MyActivation, self).__init__(**kwargs)
        # Signal that the layer is safe for mask propagation
        self.supports_masking = True

    def call(self, inputs):
        return tf.nn.relu(inputs)

これで、マスク生成レイヤー(Embedding など)とマスク消費レイヤー(LSTM など)間でこのカスタムレイヤーの使用が可能となり、マスク消費レイヤーまで届くようにマスクを渡します。

inputs = keras.Input(shape=(None,), dtype="int32")
x = layers.Embedding(input_dim=5000, output_dim=16, mask_zero=True)(inputs)
x = MyActivation()(x)  # Will pass the mask along
print("Mask found:", x._keras_mask)
outputs = layers.LSTM(32)(x)  # Will receive the mask

model = keras.Model(inputs, outputs)
Mask found: KerasTensor(type_spec=TensorSpec(shape=(None, None), dtype=tf.bool, name=None), name='Placeholder_1:0')

マスク情報が必要なレイヤーを書く

一部のレイヤーはマスクコンシューマです。それらは callmask 引数を受け取り、それを使って特定の時間ステップをスキップするかどうかを判断します。

そのようなレイヤーを書くには、単純に call シグネチャに mask=None 引数を追加します。入力に関連付けられたマスクは、それが利用可能な時にいつでもレイヤーに渡されます。

以下に簡単な例を示します。これは入力シーケンスの時間次元(軸 1)のソフトマックスを計算し、マスクされたタイムステップを破棄するレイヤーです。

class TemporalSoftmax(keras.layers.Layer):
    def call(self, inputs, mask=None):
        broadcast_float_mask = tf.expand_dims(tf.cast(mask, "float32"), -1)
        inputs_exp = tf.exp(inputs) * broadcast_float_mask
        inputs_sum = tf.reduce_sum(
            inputs_exp * broadcast_float_mask, axis=-1, keepdims=True
        )
        return inputs_exp / inputs_sum


inputs = keras.Input(shape=(None,), dtype="int32")
x = layers.Embedding(input_dim=10, output_dim=32, mask_zero=True)(inputs)
x = layers.Dense(1)(x)
outputs = TemporalSoftmax()(x)

model = keras.Model(inputs, outputs)
y = model(np.random.randint(0, 10, size=(32, 100)), np.random.random((32, 100, 1)))

要約

Keras のパディングとマスキングについて知っておくべきことはこれだけです。以下に要約します。

  • 「マスキング」とは、シーケンス入力の特定の時間ステップをスキップしたり無視したりするタイミングをレイヤーが分かるようにする方法です。
  • 一部のレイヤーはマスクジェネレーターです : Embedding は入力値からマスクを生成することができ(mask_zero=True の場合)、Masking レイヤーも同様に生成することができます。
  • 一部のレイヤーはマスクコンシューマです : これらは mask 引数を __call__ メソッドで公開します。RNN レイヤーはこれに該当します。
  • Functional API および Sequential API では、マスク情報は自動的に伝搬されます。
  • レイヤーをスタンドアロンで使用する場合には、mask 引数をレイヤーに手動で渡すことができます。
  • 現在のマスクを変更するレイヤー、新しいマスクを生成するレイヤー、入力に関連付けられたマスクを消費するレイヤーを簡単に書くことができます。