options for implementing conv max pool

max_pool.py

def max_pool_forward_naive(x, pool_param):
  """A naive implementation of the forward pass for a max-pooling layer.

  Inputs:
  - x: Input data, of shape (N, C, H, W)
  - pool_param: dictionary with the following keys:
    - 'pool_height': The height of each pooling region
    - 'pool_width': The width of each pooling region
    - 'stride': The distance between adjacent pooling regions

  No padding is necessary here, eg you can assume:
    - (H - pool_height) % stride == 0
    - (W - pool_width) % stride == 0

  Returns a tuple of:
  - out: Output data, of shape (N, C, H', W') where H' and W' are given by
    H' = 1 + (H - pool_height) / stride
    W' = 1 + (W - pool_width) / stride
  - cache: (x, pool_param)
  """
  out = None
  ###########################################################################
  # TODO: Implement the max-pooling forward pass                            #
  ###########################################################################
  # *****START OF YOUR CODE (DO NOT DELETE/MODIFY THIS LINE)*****
  pool_height = pool_param['pool_height']
  pool_width = pool_param['pool_width'] 
  stride = pool_param['stride']

  N, C, H, W = x.shape

  # prepare output tensor
  H_out = int(1 + (H - pool_height) / stride)
  W_out = int(1 + (W - pool_width) / stride)
  out = np.zeros((N, C, H_out, W_out))

  # generate columnwise store of x values for each neuron
  x_col = np.zeros((C*pool_height*pool_width, H_out*W_out))

  for i in range(N):  # loop through each image
      out_col = np.zeros((C, H_out*W_out))
      neuron_idx = 0  # loop over each neuron (= pooled segment)
      for j in range(0, 1+(H-pool_height), stride):
          for k in range(0, 1+(W-pool_width), stride):
              x_pool = x[i,:,j:j+pool_height,k:k+pool_width].reshape(
                C, pool_height*pool_width)
              out_col[:, neuron_idx] = x_pool.max(axis=1)
              neuron_idx += 1

      # update output for a batch example
      out[i] = out_col.reshape(C, H_out, W_out) 

  cache = (x, pool_param)
  return out, cache