Wrap op params for theano.gpuarray.pool.GpuPool.

theano/gpuarray/pool.c

@@ -217,8 +217,8 @@ KERNEL void ave_pool3d_kernel(const ga_size nthreads,
 
 // output shape for a given input padded shape, window shape and stride
 // We use ssize_t in the max since this is done to avoid negative results.
-#define OUTPUT_DIMS(in_dim, ws, st)                       \
-  (IGNORE_BORDER ? (in_dim - ws)/st + 1 :                 \
+#define OUTPUT_DIMS(in_dim, ws, st, ignore_border)        \
+  (ignore_border ? (in_dim - ws)/st + 1 :                 \
    (st > ws ? (in_dim - 1)/st + 1 :                       \
     std::max<ssize_t>(0, (in_dim - 1 - ws + st)/st) + 1))
 
@@ -229,7 +229,10 @@ int APPLY_SPECIFIC(pool)(PyGpuArrayObject *x,
                          PyArrayObject *stride,
                          PyArrayObject *pad,
                          PyGpuArrayObject **z,
-                         PyGpuContextObject *ctx) {
+                         PARAMS_TYPE* params) {
+  bool max_pool = (params->mode == POOLING_MAX);
+  bool inc_pad = (params->mode != POOLING_AVERAGE_COUNT_EXCLUDE_PADDING);
+  bool sum_mode = (params->mode  == POOLING_SUM);
   if (!GpuArray_IS_C_CONTIGUOUS(&x->ga))
     {
       PyErr_Format(PyExc_ValueError,
@@ -253,19 +256,19 @@ int APPLY_SPECIFIC(pool)(PyGpuArrayObject *x,
     w[i] = *((npy_int64*)PyArray_GETPTR1(ws, i));
     s[i] = *((npy_int64*)PyArray_GETPTR1(stride, i));
     p[i] = *((npy_int64*)PyArray_GETPTR1(pad, i));
-    z_dims[2 + i] = OUTPUT_DIMS(x_dims[2 + i] + 2*p[i], w[i], s[i]);
+    z_dims[2 + i] = OUTPUT_DIMS(x_dims[2 + i] + 2*p[i], w[i], s[i], params->ignore_border);
     if (p[i] > 0) {
       nonzero_padding = 1;
     }
   }
-  if (!IGNORE_BORDER && nonzero_padding) {
+  if (!params->ignore_border && nonzero_padding) {
     PyErr_SetString(PyExc_ValueError,
                     "GpuPool: padding works only with ignore_border=True");
     return 1;
   }
 
   if (theano_prep_output(z, PyGpuArray_NDIM(x), z_dims,
-                         x->ga.typecode, GA_C_ORDER, ctx) != 0)
+                         x->ga.typecode, GA_C_ORDER, params->context) != 0)
     {
       PyErr_SetString(PyExc_RuntimeError,
                       "GpuPool: failed to allocate memory");
@@ -277,7 +280,7 @@ int APPLY_SPECIFIC(pool)(PyGpuArrayObject *x,
 
     if (ndims == 2) {
       size_t num_kernels = z_dims[0] * z_dims[1] * z_dims[2] * z_dims[3];
-      if (MAX_POOL) {
+      if (max_pool) {
         err = max_pool2d_kernel_scall(1, &num_kernels, 0, num_kernels,
                                       z_dims[0], z_dims[1], z_dims[2], z_dims[3],
                                       x_dims[2], x_dims[3],
@@ -295,7 +298,7 @@ int APPLY_SPECIFIC(pool)(PyGpuArrayObject *x,
                                       x_dims[2], x_dims[3],
                                       x->ga.data, x->ga.offset,
                                       w[0], w[1], s[0], s[1], p[0], p[1],
-                                      INC_PAD, SUM_MODE,
+                                      inc_pad, sum_mode,
                                       (*z)->ga.data, (*z)->ga.offset);
         if (err != GA_NO_ERROR) {
           PyErr_Format(PyExc_RuntimeError,
@@ -307,7 +310,7 @@ int APPLY_SPECIFIC(pool)(PyGpuArrayObject *x,
     }
     else if (ndims == 3) {
       size_t num_kernels = z_dims[0] * z_dims[1] * z_dims[2] * z_dims[3] * z_dims[4];
-      if (MAX_POOL) {
+      if (max_pool) {
         err = max_pool3d_kernel_scall(1, &num_kernels, 0, num_kernels,
                                       z_dims[0], z_dims[1], z_dims[2], z_dims[3], z_dims[4],
                                       x_dims[2], x_dims[3], x_dims[4],
@@ -326,7 +329,7 @@ int APPLY_SPECIFIC(pool)(PyGpuArrayObject *x,
                                       x->ga.data, x->ga.offset,
                                       w[0], w[1], w[2], s[0], s[1], s[2],
                                       p[0], p[1], p[2],
-                                      INC_PAD, SUM_MODE,
+                                      inc_pad, sum_mode,
                                       (*z)->ga.data, (*z)->ga.offset);
         if (err != GA_NO_ERROR) {
           PyErr_Format(PyExc_RuntimeError,

theano/gpuarray/pool.py

@@ -3,9 +3,12 @@ import os.path
 
 import theano
 from theano import Apply
+from theano.gof import ParamsType
+from theano.scalar import bool as bool_t
 from theano.tensor.basic import as_tensor_variable
-from theano.tensor.signal.pool import Pool
+from theano.tensor.signal.pool import Pool, PoolingMode_t
 
+from .type import gpu_context_type
 from .basic_ops import (CGpuKernelBase, infer_context_name,
                         as_gpuarray_variable, gpu_contiguous)
 
@@ -22,6 +25,9 @@ class GpuPool(CGpuKernelBase):
 
     """
     __props__ = ('ignore_border', 'mode', 'ndim')
+    params_type = ParamsType(ignore_border=bool_t,
+                             mode=PoolingMode_t,
+                             context=gpu_context_type)
 
     def __init__(self, ignore_border, mode='max', ndim=2):
         self.ndim = ndim
@@ -31,9 +37,12 @@ class GpuPool(CGpuKernelBase):
         self.mode = mode
         CGpuKernelBase.__init__(self, ['pool.c'],
                                 'APPLY_SPECIFIC(pool)')
-        assert mode in ('max', 'sum', 'average_inc_pad', 'average_exc_pad')
+        assert PoolingMode_t.has_alias(self.mode)
         assert self.ndim in [2, 3]
 
+    def get_params(self, node):
+        return self.params_type.get_params(self, context=node.inputs[0].type.context)
+
     def c_headers(self):
         return ['gpuarray_api.h', 'gpuarray_helper.h', 'numpy_compat.h']
 
@@ -74,16 +83,6 @@ class GpuPool(CGpuKernelBase):
 
         return Apply(self, [inp, ws, stride, pad], [inp.type()])
 
-    def get_op_params(self):
-        ignore_border = int(self.ignore_border)
-        max_pool = int(self.mode == 'max')
-        inc_pad = int(self.mode != 'average_exc_pad')
-        sum_mode = int(self.mode == 'sum')
-        return [('IGNORE_BORDER', ignore_border),
-                ('INC_PAD', inc_pad),
-                ('MAX_POOL', max_pool),
-                ('SUM_MODE', sum_mode)]
-
     def infer_shape(self, node, in_shapes):
         ws, stride, pad = [node.inputs[1], node.inputs[2], node.inputs[3]]
         shp = Pool.out_shape(in_shapes[0], ws, self.ignore_border, stride,

theano/tensor/signal/pool.py

@@ -14,7 +14,7 @@ from six.moves import xrange
 import six.moves.builtins as builtins
 import theano
 from theano import gof, OpenMPOp, tensor, Variable, Apply
-from theano.gof.params_type import ParamsType
+from theano.gof import ParamsType, EnumList
 from theano.gradient import DisconnectedType
 from theano.scalar import bool as bool_t
 
@@ -258,6 +258,16 @@ def pool_3d(input, ws=None, ignore_border=None, stride=None, pad=(0, 0, 0),
     return output
 
 
+# NB: This enum type is currently used in gpuarray/pool.py.
+# It may be used later as op param in this current file.
+# Enum name and constants names are inspired from cuDNN type `cudnnPoolingMode_t`
+# (cf. `theano/gpuarray/cudnn_defs.py`).
+PoolingMode_t = EnumList(('POOLING_MAX', 'max'),
+                         ('POOLING_SUM', 'sum'),
+                         ('POOLING_AVERAGE_COUNT_INCLUDE_PADDING', 'average_inc_pad'),
+                         ('POOLING_AVERAGE_COUNT_EXCLUDE_PADDING', 'average_exc_pad'))
+
+
 class Pool(OpenMPOp):
     """
     sum or average over different patches.