title: How to optimize this MaxPool2d implementation
body: <p>I made some implementations of MaxPool2d(Running correctly, comparing with a pytorch). When testing this on a mnist dataset, this function(updateOutput) takes a very long time to complete. How to optimize this code using numpy?</p>
  
  <pre><code>class MaxPool2d(Module):
      def __init__(self, kernel_size):
          super(MaxPool2d, self).__init__()
          self.kernel_size = kernel_size
          self.gradInput = None
  
      def updateOutput(self, input):
          #print("MaxPool updateOutput")
          #start_time = time.time()
          kernel = self.kernel_size
          poolH = input.shape[2] // kernel
          poolW = input.shape[3] // kernel
          self.output = np.zeros((input.shape[0], 
                                  input.shape[1], 
                                  poolH,
                                  poolW))
          self.index = np.zeros((input.shape[0],
                                      input.shape[1],
                                      poolH,
                                      poolW,
                                      2), 
                                      dtype='int32')
  
          for i in range(input.shape[0]):
              for j in range(input.shape[1]):
                  for k in range(0, input.shape[2] - kernel+1, kernel):
                      for m in range(0, input.shape[3] - kernel+1, kernel):
                          M = input[i, j, k : k+kernel, m : m+kernel]
                          self.output[i, j, k // kernel, m // kernel] = M.max()
                          self.index[i, j, k // kernel, m // kernel] = np.array(np.unravel_index(M.argmax(), M.shape)) + np.array((k, m))
  
          #print(f"time: {time.time() - start_time:.3f}s")
          return self.output
  </code></pre>
  
  <p>input shape = (batch_size, n_input_channels, h, w)</p>
  
  <p>output shape = (batch_size, n_output_channels, h // kern_size, w // kern_size)</p>