Implement max/mean pooling(with stride) with numpy
Here's a pure numpy implementation using stride_tricks:
import numpy as npfrom numpy.lib.stride_tricks import as_strideddef pool2d(A, kernel_size, stride, padding=0, pool_mode='max'):'''2D PoolingParameters: A: input 2D array kernel_size: int, the size of the window over which we take pool stride: int, the stride of the window padding: int, implicit zero paddings on both sides of the input pool_mode: string, 'max' or 'avg''''# PaddingA = np.pad(A, padding, mode='constant')# Window view of Aoutput_shape = ((A.shape[0] - kernel_size) // stride + 1, (A.shape[1] - kernel_size) // stride + 1)shape_w = (output_shape[0], output_shape[1], kernel_size, kernel_size)strides_w = (stride*A.strides[0], stride*A.strides[1], A.strides[0], A.strides[1])A_w = as_strided(A, shape_w, strides_w)# Return the result of poolingif pool_mode == 'max': return A_w.max(axis=(2, 3))elif pool_mode == 'avg': return A_w.mean(axis=(2, 3))
Example:
>>> A = np.array([[1, 1, 2, 4], [5, 6, 7, 8], [3, 2, 1, 0], [1, 2, 3, 4]])>>> pool2d(A, kernel_size=2, stride=2, padding=0, pool_mode='max')array([[6, 8], [3, 4]])