def _test_max_pool_backward(input_shape, kernel_size, stride):
np.random.seed(0)
torch.manual_seed(0)
padding = (kernel_size - 1) // 2
input = np.random.random(input_shape).astype(np.float32) * 20
layer = MaxPoolLayer(kernel_size, stride)
torch_layer = nn.MaxPool2d(kernel_size, stride, padding)
output = layer.forward(input)
out_grad = layer.backward(2 * np.ones_like(output) / output.size)
torch_input = utils.from_numpy(input).requires_grad_(True)
torch_out = torch_layer(torch_input)
(2 * torch_out.mean()).backward()
torch_out_grad = utils.to_numpy(torch_input.grad)
utils.assert_close(out_grad, torch_out_grad, atol=TOLERANCE)
def _test_max_pool_forward(input_shape, kernel_size, stride):
np.random.seed(0)
torch.manual_seed(0)
padding = (kernel_size - 1) // 2
input = np.random.random(input_shape).astype(np.float32) * 20
original_input = input.copy()
layer = MaxPoolLayer(kernel_size, stride)
torch_layer = nn.MaxPool2d(kernel_size, stride, padding)
output = layer.forward(input)
torch_data = utils.from_numpy(input)
torch_out = utils.to_numpy(torch_layer(torch_data))
output[np.abs(output) < 1e-4] = 0
torch_out[np.abs(torch_out) < 1e-4] = 0
assert np.all(input == original_input)
assert output.shape == torch_out.shape
utils.assert_close(output, torch_out, atol=TOLERANCE)