Custom pooling layer - minmax pooling - Keras - Tensorflow

Question

I want to define my custom pooling layer, instead of returning the max values like the MaxPooling layer, it would output the k maximum values and the k minimum values.

I am using Tensorflow as the backend. I need the output vector to be sorted.

I was thinking about doing that:

from keras.layers.pooling import _Pooling1D

class MinMaxPooling1D(_Pooling1D):

    def __init__(self, output_dim, **kwargs):
        self.output_dim = output_dim
        super(MinMaxPooling1D, self).__init__(**kwargs)

    def _pooling_function(self, inputs, **kwargs):
        sorted_ = tf.contrib.framework.sort(inputs, axis = -1)
        print(sorted_)
        return np.concatenate((sorted_[:,:,:self.output_dim/2], sorted_[:,:,-self.output_dim/2:]))

But then I get:

Tensor("min_max_pooling1d_1/sort/Neg_1:0", shape=(?, 1, 1, ?), dtype=float32)
ValueError: zero-dimensional arrays cannot be concatenated

The MinMaxPooling1D layer is applied to a (None, 1, 10) shape output.

I was then thinking about adding a Flatten layer before the MinMaxPooling1D one, but then there is a dimension issue:

ValueError: Input 0 is incompatible with layer min_max_pooling1d_5: expected ndim=3, found ndim=2

Answer 1

Instead of trying to use a Pooling Layer, I am using a Lambda:

def top_k(inputs, k):
  return tf.nn.top_k(inputs, k=k, sorted=True).values

def least_k(inputs, k):
    return -tf.nn.top_k(-inputs, k=k, sorted = True).values

def minmax_k(inputs, k):
    return tf.concat([least_k(inputs, k), top_k(inputs, k)], axis = -1)

model = Sequential()
...
model.add(Lambda(minmax_k, arguments={'k': R}))