Mean or max pooling with masking support in Keras
<pre><code>...
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'))
....
</code></pre>

<p>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. </p>

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

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

<p>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:</p>

<pre><code>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)
</code></pre>

python masking keras lstm pooling