title: Python KMax Pooling (MXNet)
body: <p>I'm trying to recreate the char-level CNN in <a href="https://arxiv.org/pdf/1606.01781v1.pdf" rel="nofollow">this paper</a> and am a bit stuck at the final step where I need to create a k-max pooling layer, because I am using MXNet and it does not have this.</p>
  
  <blockquote>
    <p>An important difference is also the introduction of multiple temporal
    k-max pooling layers. This allows to detect the k most important
    features in a sentence, independent of their specific position,
    preserving their relative order.</p>
  </blockquote>
  
  <p>However, MXNet does have the ability to <a href="http://mxnet.readthedocs.io/en/latest/how_to/new_op.html" rel="nofollow">add a new-op</a> which I have been trying to do like so (although getting a bit confused with the shape of the data, given filters and batch-size). </p>
  
  <p>The shape of the data coming in:</p>
  
  <pre><code>128 (min-batch) x 512 (number of filters) x 1 (height) x 125 (width)
  </code></pre>
  
  <p>The shape of the data coming out (k-max pooling, k = 7):</p>
  
  <pre><code>128 (min-batch) x 512 (number of filters) x 1 (height)  x 7 (width)
  </code></pre>
  
  <p>My idea so far ... :</p>
  
  <pre><code>class KMaxPooling(mx.operator.CustomOp):
      def forward(self, is_train, req, in_data, out_data, aux):
          # Desired (k=3):
          # in_data = np.array([1, 2, 4, 10, 5, 3])
          # out_data = [4, 10, 5]
          x = in_data[0].asnumpy()
          idx = x.argsort()[-k:]
          idx.sort(axis=0)
          y = x[idx]
  </code></pre>
  
  <p>However, I'm not sure about several things:</p>
  
  <ol>
  <li>How to test whether this works (once I have some complete code)</li>
  <li>What the dimensions should be? I'm sorting on the last dimension (axis=0)</li>
  <li>What to do for the backward() function i.e. the gradient propogation</li>
  <li>Whether this will work with GPU - I'm guessing I will have to rewrite it in C/cuda?</li>
  </ol>
  
  <p>I found this example by someone else for keras (but don't have the rep to link):</p>
  
  <pre><code>import numpy as np
  import theano.tensor as T
  from keras.layers.core import MaskedLayer
  
  class KMaxPooling(MaskedLayer):
      def __init__(self, pooling_size):
          super(MaskedLayer, self).__init__()
          self.pooling_size = pooling_size
          self.input = T.tensor3()
  
      def get_output_mask(self, train=False):
          return None
  
      def get_output(self, train=False):
          data = self.get_input(train)
          mask = self.get_input_mask(train)
  
          if mask is None:
              mask = T.sum(T.ones_like(data), axis=-1)
          mask = mask.dimshuffle(0, 1, "x")
  
          masked_data = T.switch(T.eq(mask, 0), -np.inf, data)
  
          result = masked_data[T.arange(masked_data.shape[0]).dimshuffle(0, "x", "x"),
                               T.sort(T.argsort(masked_data, axis=1)[:, -self.pooling_size:, :], axis=1),
                               T.arange(masked_data.shape[2]).dimshuffle("x", "x", 0)]
  </code></pre>