wehub-resource-sync
94 lines
3.1 KiB
Python
Executable File
94 lines
3.1 KiB
Python
Executable File
import numpy as np
|
|
from prml.nn.tensor.tensor import Tensor
|
|
from prml.nn.function import Function
|
|
from prml.nn.image.util import img2patch, patch2img
|
|
|
|
|
|
class MaxPooling2d(Function):
|
|
|
|
def __init__(self, pool_size, stride, pad):
|
|
"""
|
|
construct 2 dimensional max-pooling function
|
|
|
|
Parameters
|
|
----------
|
|
pool_size : int or tuple of ints
|
|
pooling size
|
|
stride : int or tuple of ints
|
|
stride of kernel application
|
|
pad : int or tuple of ints
|
|
padding image
|
|
"""
|
|
self.pool_size = self._check_tuple(pool_size, "pool_size")
|
|
self.stride = self._check_tuple(stride, "stride")
|
|
self.pad = self._check_tuple(pad, "pad")
|
|
self.pad = (0,) + self.pad + (0,)
|
|
|
|
def _check_tuple(self, tup, name):
|
|
if isinstance(tup, int):
|
|
tup = (tup,) * 2
|
|
if not isinstance(tup, tuple):
|
|
raise TypeError(
|
|
"Unsupported type for {}: {}".format(name, type(tup))
|
|
)
|
|
if len(tup) != 2:
|
|
raise ValueError(
|
|
"the length of {} must be 2, not {}".format(name, len(tup))
|
|
)
|
|
if not all([isinstance(n, int) for n in tup]):
|
|
raise TypeError(
|
|
"Unsuported type for {}".format(name)
|
|
)
|
|
if not all([n >= 0 for n in tup]):
|
|
raise ValueError(
|
|
"{} must be non-negative values".format(name)
|
|
)
|
|
return tup
|
|
|
|
def forward(self, x):
|
|
x = self._convert2tensor(x)
|
|
self._equal_ndim(x, 4)
|
|
self.x = x
|
|
img = np.pad(x.value, [(p,) for p in self.pad], "constant")
|
|
patch = img2patch(img, self.pool_size, self.stride)
|
|
n_batch, xlen_out, ylen_out, _, _, in_channels = patch.shape
|
|
patch = patch.reshape(n_batch, xlen_out, ylen_out, -1, in_channels)
|
|
self.shape = img.shape
|
|
self.index = patch.argmax(axis=3)
|
|
return Tensor(patch.max(axis=3), function=self)
|
|
|
|
def backward(self, delta):
|
|
delta_patch = np.zeros(delta.shape + (np.prod(self.pool_size),))
|
|
index = np.where(delta == delta) + (self.index.ravel(),)
|
|
delta_patch[index] = delta.ravel()
|
|
delta_patch = np.reshape(delta_patch, delta.shape + self.pool_size)
|
|
delta_patch = delta_patch.transpose(0, 1, 2, 4, 5, 3)
|
|
dx = patch2img(delta_patch, self.stride, self.shape)
|
|
slices = [slice(p, len_ - p) for p, len_ in zip(self.pad, self.shape)]
|
|
dx = dx[slices]
|
|
self.x.backward(dx)
|
|
|
|
|
|
def max_pooling2d(x, pool_size, stride=1, pad=0):
|
|
"""
|
|
spatial max pooling
|
|
|
|
Parameters
|
|
----------
|
|
x : (n_batch, xlen, ylen, in_channel) Tensor
|
|
input tensor
|
|
pool_size : int or tuple of ints (kx, ky)
|
|
pooling size
|
|
stride : int or tuple of ints (sx, sy)
|
|
stride of pooling application
|
|
pad : int or tuple of ints (px, py)
|
|
padding input
|
|
|
|
Returns
|
|
-------
|
|
output : (n_batch, xlen', ylen', out_channel) Tensor
|
|
max pooled image
|
|
len' = (len + p - k) // s + 1
|
|
"""
|
|
return MaxPooling2d(pool_size, stride, pad).forward(x)
|