Notes on masking and padding in tensorflow keras

1. background

Masking The function of is to inform the sequence processing layer that some time steps are lost in the input , Therefore, it should be skipped when processing data .
padding yes Masking A special form of , Among them, the Masking The step of is at the beginning or beginning of the sequence . Filling is required to encode sequence data into successive batches ： In order to fit all sequences in the batch to a given standard length , It is necessary to fill or truncate some sequences .

2. padding Example of filling sequence data

When processing sequence data , Each sample often has a different length . Consider the following example （ The text is segmented into words ）：

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

Conduct embedding After query , Data may be vectorized into integers , for example ：

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

This data is a nested list , The length of each sample is 3、5 and 6. Because the input data of the deep learning model must be a single tensor （ For example, in this example, the shape is (batch_size, 6, vocab_size)）, Samples shorter than the longest entry need to be filled with placeholder values （ perhaps , Long samples can also be truncated before filling short samples ）.

Keras A utility function is provided to truncate and fill Python list , Make it the same length ：tf.keras.preprocessing.sequence.pad_sequences.

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

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

3. cover （masking ）

Now that all samples now have a uniform length , Then the model must be informed , Some parts of the data are actually populated , Should be ignored . This mechanism is cover .
stay Keras There are three ways to introduce input masks into the model ：

3.1 Add one keras.layers.Masking layer .

import tensorflow as tf
import tensorflow.keras.layers as layers
raw_inputs = [
    [711, 632, 71],
    [73, 8, 3215, 55, 927],
    [83, 91, 1, 645, 1253, 927],
]
inputs = tf.keras.preprocessing.sequence.pad_sequences(
    raw_inputs, padding="post"
)
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(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)

3.2 Use mask_zero=True To configure a keras.layers.Embedding layer .

import tensorflow as tf
import tensorflow.keras.layers as layers
raw_inputs = [
    [711, 632, 71],
    [73, 8, 3215, 55, 927],
    [83, 91, 1, 645, 1253, 927],
]
inputs = tf.keras.preprocessing.sequence.pad_sequences(
    raw_inputs, padding="post"
)
embedding = layers.Embedding(input_dim=5000, output_dim=16, mask_zero=True)
masked_output = embedding(inputs)

print(masked_output._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)

3.3. On call support mask Layer of parameters （ Such as RNN layer ） when , Pass this parameter manually .

import tensorflow as tf
import tensorflow.keras.layers as layers
raw_inputs = [
    [711, 632, 71],
    [73, 8, 3215, 55, 927],
    [83, 91, 1, 645, 1253, 927],
]
inputs = tf.keras.preprocessing.sequence.pad_sequences(
    raw_inputs, padding="post"
)

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(4, return_sequences=True)

    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()
print(layer(inputs))

In fact, you can not call the above code yourself On the upper floor compute_mask() obtain mask Pass to LSTM, Because if you were Embedding Set up inside. mask_zero=True. that LSTM Automatically called compute_mask Methodical . Only some deep network or activation functions are not implemented , If so, you can only get it by yourself mask 了 .

    import tensorflow as tf
    import tensorflow.keras.layers as layers
    raw_inputs = [
        [711, 632, 71],
        [73, 8, 3215, 55, 927],
        [83, 91, 1, 645, 1253, 927],
    ]

    inputs = tf.keras.preprocessing.sequence.pad_sequences(
        raw_inputs, padding="post"
    )
    print(inputs)
    x = layers.Embedding(input_dim=5000, output_dim=5, mask_zero=True)(inputs)
    x =  tf.nn.relu(x)
    x = layers.LSTM(4, return_sequences=True)(x)
    print(x)

tf.Tensor(
[[[-0.00853076 -0.0027577   0.00382648 -0.00217819]
  [-0.00701238 -0.00343913  0.0032082  -0.0011041 ]
  [-0.00747839 -0.00233405  0.00432012 -0.00373881]
  [-0.0137679  -0.00557745  0.00911791 -0.00396027]
  [-0.01824936 -0.00863818  0.01233672 -0.00297129]
  [-0.02145157 -0.01128308  0.01439027 -0.00143419]]
 [[-0.0079427  -0.00146676  0.00085193 -0.00298749]
  [-0.01137889 -0.00377863  0.00144354 -0.00350932]
  [-0.0115746  -0.00302937 -0.00108732 -0.001251  ]
  [-0.01822513 -0.00597077 -0.00068765  0.00095661]
  [-0.02203784 -0.00890452  0.00246888  0.00270562]
  [-0.02512554 -0.01110168  0.00644199  0.00187145]]
 [[ 0.          0.          0.          0.        ]
  [-0.0061018  -0.00086632 -0.00171513 -0.0001184 ]
  [-0.01114718 -0.00349189  0.00514382 -0.00568866]
  [-0.01592513 -0.00512891  0.00192113  0.00023104]
  [-0.01553659 -0.00651678  0.00391222 -0.00226363]
  [-0.01969725 -0.00938978  0.00616115  0.00070853]]], shape=(3, 6, 4), dtype=float32)

You can see masking after Relu And then it failed . haven't you? mask Pass on to lstm layer
So we should make some changes , Here are two ways ,
Method 1
encapsulation Relu Method , Set properties supports_masking=True, The code is as follows ：

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)

Method 2
utilize _keras_mask Method to get embedding Layer of mask Pass to lstm. The code is as follows ：

import tensorflow as tf
import tensorflow.keras.layers as layers
raw_inputs = [
    [711, 632, 71],
    [73, 8, 3215, 55, 927],
    [83, 91, 1, 645, 1253, 927],
]

inputs = tf.keras.preprocessing.sequence.pad_sequences(
    raw_inputs, padding="post"
)
print(inputs)
x = layers.Embedding(input_dim=5000, output_dim=5, mask_zero=True)(inputs)
mask = x._keras_mask
x =  tf.nn.relu(x)
x = layers.LSTM(4, return_sequences=True)(x, mask=mask)
print(x)

tf.Tensor(
[[[-1.11592971e-02  1.13475965e-02  1.22134504e-03  4.27628960e-03]
  [-1.17472848e-02  1.37128066e-02  3.33235669e-03  4.49653203e-03]
  [-1.18236477e-02  1.50292199e-02 -1.04225962e-03  6.14795834e-03]
  [-1.18236477e-02  1.50292199e-02 -1.04225962e-03  6.14795834e-03]
  [-1.18236477e-02  1.50292199e-02 -1.04225962e-03  6.14795834e-03]
  [-1.18236477e-02  1.50292199e-02 -1.04225962e-03  6.14795834e-03]]
 [[-3.03814257e-03  3.01474496e-03 -2.19832035e-03  1.65364845e-03]
  [-3.02985404e-03 -2.74110964e-04 -8.03735293e-03  7.56251905e-03]
  [-4.05789400e-03  6.13932591e-03 -7.65458029e-03  9.17674601e-03]
  [-1.01635903e-02  1.46681387e-02 -2.00540805e-03  1.00620193e-02]
  [-9.56059992e-03  1.74404141e-02 -2.37406185e-03  1.30043132e-02]
  [-9.56059992e-03  1.74404141e-02 -2.37406185e-03  1.30043132e-02]]

4 Write layers that require mask information

Some layers are mask consumers ： They will be in call Accept in mask Parameters , And use this parameter to decide whether to skip some time steps .

To write such a layer , All you need to do is call Add a... To the signature mask=None Parameters . The mask associated with the input is passed to your layer whenever it is available .

Here is a simple example ： The layer in the example is in the time dimension of the input sequence （ Axis 1） Count up Softmax, At the same time, discard the time steps of the cover .

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 * broadcast_float_mask, axis=1, keepdims=True)
        returcn 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)))