Mean or max pooling with masking support in Keras

Question

...
print('Build model...')
model = Sequential()
model.add(Embedding(max_features, 128))
model.add(LSTM(size, return_sequences=True, dropout_W=0.2 dropout_U=0.2)) 
model.add(GlobalAveragePooling1D())
model.add(Dense(1))
model.add(Activation('sigmoid'))
....

I need to be able to take the mean or max of the vectors for all time steps in a sample after LSTM layer before giving this mean or max vector to the dense layer in Keras.

I think timedistributedmerge was able to do this but it was deprecated. Using return_sequences=True I can obtain the vectors for all time steps in a sample after the LSTM layer. However, GlobalAveragePooling1D() is not compatible with masking and it considers all time steps whereas I need only the non-masked time steps.

I saw posts recommending the Lambda layer but these also do not take masking into account. Any help would be appreciated.

score 4 · Answer 1 · answered Aug 17 '17 at 08:52

Jacoxu's answer is right. But if you are using a tensorflow backend for keras, the Tensor type doesn't support dimshuffle function, try this instead.

def call(self, x, mask=None):
    if mask is not None:
        # mask (batch, time)
        mask = K.cast(mask, K.floatx())
        # mask (batch, x_dim, time)
        mask = K.repeat(mask, x.shape[-1])
        # mask (batch, time, x_dim)
        mask = tf.transpose(mask, [0,2,1])
        x = x * mask
    return K.sum(x, axis=1) / K.sum(mask, axis=1)

score 3 · Answer 2 · answered Sep 16 '16 at 14:35

3

Since average pooling is only doing a mean over one axis, you just need to correct the number of elements in the mean since loss masking is handled at the end, not here. You can do this probably with something like this:

class GlobalAveragePooling1DMasked(GlobalAveragePooling1D):
    def call(self, x, mask=None):
        if mask != None:
            return K.sum(x, axis=1) / K.sum(mask, axis=1)
        else:
            return super().call(x)

answered Sep 16 '16 at 14:35

nemo

55,207
13
135
135

2

Note that you can't be certain in Keras that masked values in `x` are equal to zero! Therefore, this implementation would give wrong results. – pir Mar 04 '17 at 20:37
The mask itself, introduced by the `Masking` and `Embedding` layer, is binary. Of course you can always have layers that implement `compute_mask` differently, but this is not happening in keras itself as far as I can see. – nemo Mar 05 '17 at 03:47
Yes, the mask is binary. That is not what I'm saying :) – pir Mar 05 '17 at 04:25
1

But that's the only assumption I make in the example above. I don't understand what you are getting at, sorry. Edit: Nevermind, it's the sum operation of matrix `x` that makes the assumption of masked values being zero. You are right. This *could* probably be fixed by indexing `x` by `mask` but I have to test this first. – nemo Mar 05 '17 at 15:19

score 3 · Answer 3 · answered Mar 22 '17 at 05:54

In order to make the masked values in x be equal to zero, you can do this:

class MeanPool(Layer):
def __init__(self, **kwargs):
    self.supports_masking = True
    super(MeanPool, self).__init__(**kwargs)

def compute_mask(self, input, input_mask=None):
    # do not pass the mask to the next layers
    return None

def call(self, x, mask=None):
    if mask is not None:
        # mask (batch, time)
        mask = K.cast(mask, K.floatx())
        # mask (batch, time, 'x')
        mask = mask.dimshuffle(0, 1, 'x')
        # to make the masked values in x be equal to zero
        x = x * mask
    return K.sum(x, axis=1) / K.sum(mask, axis=1)

def get_output_shape_for(self, input_shape):
    # remove temporal dimension
    return input_shape[0], input_shape[2]

score 3 · Answer 4 · answered May 02 '18 at 23:10

This is how I did it on Keras 2 (borrowing from all of the answers, and fixing the dimensions):

class MeanPool(Layer):
  def __init__(self, **kwargs):
      self.supports_masking = True
      super(MeanPool, self).__init__(**kwargs)

  def compute_mask(self, input, input_mask=None):
      # do not pass the mask to the next layers
      return None

  def call(self, x, mask=None):
      if mask is not None:
          # mask (batch, time)
          mask = K.cast(mask, K.floatx())
          # mask (batch, x_dim, time)
          mask = K.repeat(mask, x.shape[-1])
          # mask (batch, time, x_dim)
          mask = tf.transpose(mask, [0,2,1])
          x = x * mask
      return K.sum(x, axis=1) / K.sum(mask, axis=1)

  def compute_output_shape(self, input_shape):
      # remove temporal dimension
      return (input_shape[0], input_shape[2])

Mean or max pooling with masking support in Keras

4 Answers4

Linked