Merge pull request #3486 from niboshi/average-pooling-pad-value

Support skipping padded regions in `F.average_pooling_nd`

chainer/functions/pooling/average_pooling_2d.py

@@ -210,5 +210,12 @@ def average_pooling_2d(x, ksize, stride=None, pad=0):
        This function currently does not support ``cover_all`` mode as
        :func:`max_pooling_2d`. Average pooling runs in non-cover-all mode.
 
+    .. note::
+
+       The values in the padded region is treated as 0, leading the averages
+       biased towards zero.
+       To obtain unbiased averages, use :func:`average_pooling_nd` with
+       ``pad_value=None``.
+
     """
     return AveragePooling2D(ksize, stride, pad, False).apply((x,))[0]

chainer/functions/pooling/average_pooling_nd.py

@@ -10,9 +10,27 @@
 from chainer.functions.pooling import average_pooling_nd_kernel
 from chainer.functions.pooling import pooling_nd
 from chainer import utils
+from chainer.utils import conv
 from chainer.utils import conv_nd
 
 
+def _get_conv_slices(
+        size, k, s, p, cover_all=False, d=1, include_pad=True, dtype='l'):
+    """Returns the patch slices.
+
+    Returns:
+        A tuple of two 1-D :class:`numpy.ndarrays`\\ s.
+        Each represents starting and ending indices of the patches.
+    """
+    n = conv.get_conv_outsize(size, k, s, p, cover_all, d)
+    starts = -p + numpy.arange(n, dtype=dtype) * s
+    ends = starts + k
+    if not include_pad:
+        starts = numpy.maximum(starts, 0)
+        ends = numpy.minimum(ends, size)
+    return starts, ends
+
+
 class AveragePoolingND(pooling_nd._PoolingND):
 
     """Average pooling over a set of N-dimensional planes.
@@ -23,7 +41,13 @@ class AveragePoolingND(pooling_nd._PoolingND):
 
     """
 
-    def __init__(self, ndim, ksize, stride=None, pad=0, cover_all=False):
+    def __init__(
+            self, ndim, ksize, stride=None, pad=0, cover_all=False,
+            pad_value=0):
+        if not (pad_value is None or pad_value == 0):
+            raise ValueError(
+                'pad_value must be either 0 or None, not {}.'.format(
+                    pad_value))
         utils.experimental('chainer.functions.pooling.AveragePoolingND')
 
         # TODO(takagi) Support cover_all mode.
@@ -33,55 +57,95 @@ def __init__(self, ndim, ksize, stride=None, pad=0, cover_all=False):
         super(AveragePoolingND, self).__init__(
             ndim, ksize, stride=stride, pad=pad, cover_all=cover_all)
 
-    def forward_cpu(self, x):
-        self._in_shape = x[0].shape
-        self._in_dtype = x[0].dtype
+        self.pad_value = pad_value
+
+    def _get_pooling_width(self, xp, dims, dtype):
+        width = None
+        for d, k, s, p in six.moves.zip(
+                dims, self.ksize, self.stride, self.pad):
+            starts, ends = _get_conv_slices(
+                d, k, s, p, cover_all=self.cover_all, include_pad=False,
+                dtype=dtype)
+            w = ends - starts
+            if width is None:
+                width = w
+            else:
+                width = numpy.tensordot(width[..., None], w[None, ...], axes=1)
+        if xp is not numpy:
+            width = cuda.cupy.array(width)
+        return width
+
+    def forward_cpu(self, inputs):
+        x, = inputs
+        self._in_shape = x.shape
+        self._in_dtype = x.dtype
 
         col = conv_nd.im2col_nd_cpu(
-            x[0], self.ksize, self.stride, self.pad, cover_all=self.cover_all)
+            x, self.ksize, self.stride, self.pad, cover_all=self.cover_all)
 
         # mean along (_, _, k_1, k_2, ..., k_N, _, ..., _)
         y_axis = tuple(six.moves.range(2, 2 + len(self.ksize)))
-        y = col.mean(axis=y_axis)
+        if self.pad_value is None:
+            dims = x.shape[2:]
+            width = self._get_pooling_width(numpy, dims, x.dtype)
+            y = col.sum(axis=y_axis) / width
+        else:
+            assert self.pad_value == 0
+            y = col.mean(axis=y_axis)
+
         return y,
 
-    def forward_gpu(self, x):
+    def forward_gpu(self, inputs):
         if chainer.should_use_cudnn('>=auto') and 2 <= self.ndim <= 3:
             # With cuDNN v3 or greater, use cuDNN implementation for inputs
             # with spatial dimensions of two or more.
             self.retain_inputs((0,))
-            return super(AveragePoolingND, self).forward_gpu(x)
-
-        self._in_shape = x[0].shape
-        self._in_dtype = x[0].dtype
-
-        n, c = x[0].shape[:2]
-        dims = x[0].shape[2:]
-        ys = tuple(conv_nd.get_conv_outsize(d, k, s, p,
-                                            cover_all=self.cover_all)
-                   for (d, k, s, p) in six.moves.zip(
-                       dims, self.ksize, self.stride, self.pad))
+            return super(AveragePoolingND, self).forward_gpu(inputs)
+
+        x, = inputs
+        self._in_shape = x.shape
+        self._in_dtype = x.dtype
+
+        n, c = x.shape[:2]
+        idims = x.shape[2:]
+        odims = tuple(
+            conv.get_conv_outsize(d, k, s, p, cover_all=self.cover_all)
+            for (d, k, s, p) in six.moves.zip(
+                idims, self.ksize, self.stride, self.pad))
         # (n, c, y_1, y_2, ..., y_N)
-        y_shape = (n, c) + ys
-        y = cuda.cupy.empty(y_shape, dtype=x[0].dtype)
-        coeff = 1. / functools.reduce(operator.mul, self.ksize)
+        y_shape = (n, c) + odims
+        y = cuda.cupy.empty(y_shape, dtype=x.dtype)
+        if self.pad_value is None:
+            coeff = self._get_pooling_width(cuda.cupy, idims, x.dtype)
+            coeff = cuda.cupy.reciprocal(coeff, out=coeff)
+        else:
+            assert self.pad_value == 0
+            coeff = 1. / functools.reduce(operator.mul, self.ksize)
 
         in_params, out_params, operation, name = \
             average_pooling_nd_kernel.AveragePoolingNDKernelForward.generate(
                 self.ndim)
         cuda.elementwise(in_params, out_params, operation, name)(
-            x[0].reduced_view(),
-            *(dims + ys + self.ksize + self.stride + self.pad + (coeff, y)))
+            x.reduced_view(),
+            *(idims + odims + self.ksize + self.stride + self.pad
+              + (coeff, y)))
 
         return y,
 
     def backward(self, indexes, gy):
         return AveragePoolingNDGrad(self).apply(gy)
 
     def create_pool_desc(self):
+        if self.pad_value is None:
+            pooling_mode = (
+                cuda.cuda.cudnn.CUDNN_POOLING_AVERAGE_COUNT_EXCLUDE_PADDING)
+        else:
+            assert self.pad_value == 0
+            pooling_mode = (
+                cuda.cuda.cudnn.CUDNN_POOLING_AVERAGE_COUNT_INCLUDE_PADDING)
+
         return cuda.cudnn.create_pooling_descriptor(
-            self.ksize, self.stride, self.pad,
-            cuda.cuda.cudnn.CUDNN_POOLING_AVERAGE_COUNT_INCLUDE_PADDING)
+            self.ksize, self.stride, self.pad, pooling_mode)
 
 
 class AveragePoolingNDGrad(function_node.FunctionNode):
@@ -96,36 +160,48 @@ def __init__(self, apoolnd):
         if not self._used_cudnn:
             self._in_shape = apoolnd._in_shape
             self._in_dtype = apoolnd._in_dtype
+        self.pad_value = apoolnd.pad_value
         self.apoolnd = apoolnd
 
-    def forward_cpu(self, gy):
-        dims = self._in_shape[2:]
-        outs = gy[0].shape[2:]
+    def forward_cpu(self, gys):
+        gy, = gys
+        idims = self._in_shape[2:]
+        odims = gy.shape[2:]
         colon = slice(None, None, None)
-        gy_index = (colon, colon) + (None,) * len(dims)
-        gcol_reps = (1, 1) + self.ksize + (1,) * len(outs)
-        gcol = numpy.tile(gy[0][gy_index], gcol_reps)
-        gx = conv_nd.col2im_nd_cpu(gcol, self.stride, self.pad, dims)
-        gx /= functools.reduce(operator.mul, self.ksize)
+        gy_index = (colon, colon) + (None,) * len(idims)
+        gcol_reps = (1, 1) + self.ksize + (1,) * len(odims)
+        gcol = numpy.tile(gy[gy_index], gcol_reps)
+        gx = conv_nd.col2im_nd_cpu(gcol, self.stride, self.pad, idims)
+        if self.pad_value is None:
+            width = self._get_pooling_width(numpy, odims, gx.dtype)
+            numpy.divide(gx, width, out=gx)
+        else:
+            gx /= functools.reduce(operator.mul, self.ksize)
         return gx,
 
-    def forward_gpu(self, gy):
+    def forward_gpu(self, gys):
         if self._used_cudnn:
             x, = self.apoolnd.get_retained_inputs()
-            return self.apoolnd.backward_gpu((x.data,), gy)
+            return self.apoolnd.backward_gpu((x.data,), gys)
 
+        gy, = gys
         n, c = self._in_shape[:2]
-        dims = self._in_shape[2:]
-        ys = gy[0].shape[2:]
+        idims = self._in_shape[2:]
+        odims = gy.shape[2:]
         gx = cuda.cupy.empty(self._in_shape, self._in_dtype)
-        coeff = 1. / functools.reduce(operator.mul, self.ksize)
+        if self.pad_value is None:
+            coeff = self._get_pooling_width(cuda.cupy, odims, gy.dtype)
+            coeff = cuda.cupy.reciprocal(coeff, out=coeff)
+        else:
+            coeff = 1. / functools.reduce(operator.mul, self.ksize)
 
         in_params, out_params, operation, name = \
             average_pooling_nd_kernel.AveragePoolingNDKernelBackward.generate(
                 self.ndim)
         cuda.elementwise(in_params, out_params, operation, name)(
-            gy[0].reduced_view(),
-            *(dims + ys + self.ksize + self.stride + self.pad + (coeff, gx)))
+            gy.reduced_view(),
+            *(idims + odims + self.ksize + self.stride + self.pad
+              + (coeff, gx)))
 
         return gx,
 
@@ -135,7 +211,7 @@ def backward(self, indexes, grad_outputs):
             cover_all=False).apply(grad_outputs)
 
 
-def average_pooling_nd(x, ksize, stride=None, pad=0):
+def average_pooling_nd(x, ksize, stride=None, pad=0, pad_value=0):
     """N-dimensionally spatial average pooling function.
 
     .. warning::
@@ -158,6 +234,11 @@ def average_pooling_nd(x, ksize, stride=None, pad=0):
             window size.
         pad (int or tuple of ints): Spatial padding width for the input array.
             ``pad=p`` and ``pad=(p, p, ..., p)`` are equivalent.
+        pad_value (0 or None):
+            Value to fill the padded region when calculating average.
+            If ``None`` is specified, such region is ignored.
+            The default value is ``0``, therefore the averages are biased
+            towards zero.
 
     Returns:
         ~chainer.Variable: Output variable.
@@ -169,4 +250,6 @@ def average_pooling_nd(x, ksize, stride=None, pad=0):
 
     """
     ndim = len(x.shape[2:])
-    return AveragePoolingND(ndim, ksize, stride=stride, pad=pad).apply((x,))[0]
+    return AveragePoolingND(
+        ndim, ksize, stride=stride, pad=pad, pad_value=pad_value
+    ).apply((x,))[0]

chainer/functions/pooling/pooling_nd.py

@@ -85,9 +85,9 @@ def backward_gpu(self, x, gy):
         zero = numpy.array(0, dtype=oz_dtype).ctypes
         gx = cuda.cupy.empty_like(x)
         libcudnn.poolingBackward(
-            handle, pool_desc.value, one.data, y_desc.value,
-            y.data.ptr, y_desc.value, gy.data.ptr, x_desc.value,
-            x.data.ptr, zero.data, x_desc.value, gx.data.ptr)
+            handle, pool_desc.value,
+            one.data, y_desc.value, y.data.ptr, y_desc.value, gy.data.ptr,
+            x_desc.value, x.data.ptr, zero.data, x_desc.value, gx.data.ptr)
         return gx,
 
     def create_pool_desc(self):

tests/chainer_tests/functions_tests/pooling_tests/test_average_pooling_nd.py

@@ -19,6 +19,7 @@
 @testing.parameterize(*testing.product({
     'dims': [(4,), (4, 3), (4, 3, 2), (1, 1, 1, 1)],
     'dtype': [numpy.float16, numpy.float32, numpy.float64],
+    'pad_value': [None, 0],
 }))
 class TestAveragePoolingND(unittest.TestCase):
 
@@ -52,19 +53,29 @@ def check_forward(self, x_data, use_cudnn='always'):
         pad = self.pad
         x = chainer.Variable(x_data)
         with chainer.using_config('use_cudnn', use_cudnn):
-            y = functions.average_pooling_nd(x, ksize, stride, pad)
+            y = functions.average_pooling_nd(
+                x, ksize, stride, pad, self.pad_value)
         self.assertEqual(y.data.dtype, self.dtype)
         y_data = cuda.to_cpu(y.data)
 
+        def denom(idx):
+            if self.pad_value is None:
+                s = 1
+                for slic in idx:
+                    s *= slic.stop - slic.start
+                return s
+            else:
+                return functools.reduce(operator.mul, ksize)
+
         self.assertEqual(self.gy.shape, y_data.shape)
         patches = pooling_nd_helper.pooling_patches(
             dims, ksize, stride, pad, False)
         for k in six.moves.range(2):
             for c in six.moves.range(3):
                 x = self.x[k, c]
-                size = functools.reduce(operator.mul, ksize)
-                expect = numpy.array([x[idx].sum() for idx in patches])
-                expect = expect.reshape(y_data.shape[2:]) / size
+                expect = numpy.array(
+                    [x[idx].sum() / denom(idx) for idx in patches])
+                expect = expect.reshape(y_data.shape[2:])
                 testing.assert_allclose(
                     expect, y_data[k, c], **self.check_forward_options)
 
@@ -89,15 +100,20 @@ def check_forward_consistency_regression(self, x_data, use_cudnn='always'):
         if len(self.dims) != 2:
             return
 
+        if self.pad_value != 0:
+            # Not supported in average_pooling_2d
+            return
+
         ksize = self.ksize
         stride = self.stride
         pad = self.pad
 
         with chainer.using_config('use_cudnn', use_cudnn):
-            y_nd = functions.average_pooling_nd(x_data, ksize, stride=stride,
-                                                pad=pad)
-            y_2d = functions.average_pooling_2d(x_data, ksize, stride=stride,
-                                                pad=pad)
+            y_nd = functions.average_pooling_nd(
+                x_data, ksize, stride=stride, pad=pad,
+                pad_value=self.pad_value)
+            y_2d = functions.average_pooling_2d(
+                x_data, ksize, stride=stride, pad=pad)
         testing.assert_allclose(y_nd.data, y_2d.data)
 
     def test_forward_consistency_regression_cpu(self):