Wrap Op params for theano.gpuarray.dnn.GpuDnnBatchNorm.

doc/extending/extending_theano_c.txt

@@ -894,7 +894,7 @@ If you pass a function name to the ``__init__()`` method of the
         theano Types) of your inputs and outputs.
 
 *       You can sepcify the number of inputs and outputs for your op
-        by setting the `_cop_num_inputs` and `_cop_num_outputs`
+        by setting the ``_cop_num_inputs`` and ``_cop_num_outputs``
         attributes on your op.  The main function will always be
         called with that number of arguments, using NULL to fill in
         for missing values at the end.  This can be used if your op

theano/gpuarray/dnn.py

@@ -1666,13 +1666,20 @@ class GpuDnnBatchNorm(DnnBase):
     __props__ = ('mode', 'running_averages', 'inplace_running_mean',
                  'inplace_running_var', 'inplace_output')
 
+    check_input = False
+    params_type = ParamsType(mode=cudnn.cudnnBatchNormMode_t,
+                             inplace_output=bool_t,
+                             inplace_running_mean=bool_t,
+                             inplace_running_var=bool_t,
+                             handle=handle_type)
+
     def __init__(self, mode='per-activation', running_averages=False,
                  inplace_running_mean=False, inplace_running_var=False,
                  inplace_output=False):
         DnnBase.__init__(self, ['dnn_batchnorm_base.c', 'dnn_batchnorm.c'],
                          'dnn_batchnorm_op')
 
-        assert (mode in ('per-activation', 'spatial'))
+        assert cudnn.cudnnBatchNormMode_t.has_alias(mode)
         self.mode = mode
         self.running_averages = running_averages
         self.inplace_output = inplace_output
@@ -1700,24 +1707,12 @@ class GpuDnnBatchNorm(DnnBase):
             self.inplace_output = False
             self.destroy_map = {}
 
-    def get_op_params(self):
-        params = []
-        if self.inplace_output:
-            params.append(('INPLACE_OUTPUT', '1'))
-        if self.running_averages:
-            params.append(('RUNNING_AVERAGES', '1'))
-            if self.inplace_running_mean:
-                params.append(('INPLACE_RUNNING_MEAN', '1'))
-            if self.inplace_running_var:
-                params.append(('INPLACE_RUNNING_VAR', '1'))
-        params.append(('MODE', ("CUDNN_BATCHNORM_SPATIAL"
-                                if self.mode == "spatial"
-                                else "CUDNN_BATCHNORM_PER_ACTIVATION")))
-        return params
-
     def infer_shape(self, node, shape):
         return [shape[0]] + [shape[1]] * (len(node.outputs) - 1)
 
+    _cop_num_inputs = 7
+    _cop_num_outputs = 5
+
     def make_node(self, x, scale, bias, epsilon=1e-4,
                   running_average_factor=0.1,
                   running_mean=None, running_var=None):

theano/gpuarray/dnn_batchnorm.c

@@ -3,18 +3,17 @@
 int dnn_batchnorm_op(PyGpuArrayObject *inp, PyGpuArrayObject *scale,
                      PyGpuArrayObject *bias, npy_float64 epsilon,
                      npy_float64 running_average_factor,
-#ifdef RUNNING_AVERAGES
-                     PyGpuArrayObject *in_running_mean,
-                     PyGpuArrayObject *in_running_var,
-#endif
+                     PyGpuArrayObject *in_running_mean, // may be NULL
+                     PyGpuArrayObject *in_running_var, // may be NULL
                      PyGpuArrayObject **outp,
                      PyGpuArrayObject **x_mean,
                      PyGpuArrayObject **x_invstd,
-#ifdef RUNNING_AVERAGES
-                     PyGpuArrayObject **out_running_mean,
-                     PyGpuArrayObject **out_running_var,
-#endif
-                     cudnnHandle_t _handle) {
+                     PyGpuArrayObject **out_running_mean, // may be NULL
+                     PyGpuArrayObject **out_running_var, // may be NULL
+                     PARAMS_TYPE* params) {
+  /* Note: based on Python code, in_running_mean, in_running_var, out_running_mean and out_running_var
+  are together NULL (or not NULL) at same time, so we just need to check only one of them. */
+  bool running_averages = (in_running_mean != NULL);
   PyGpuContextObject *c = inp->context;
 
   if (c_set_tensorNd(inp, bn_input) != 0)
@@ -27,14 +26,14 @@ int dnn_batchnorm_op(PyGpuArrayObject *inp, PyGpuArrayObject *scale,
     return 1;
   }
 
-#ifdef INPLACE_OUTPUT
-  Py_XDECREF(*outp);
-  *outp = inp;
-  Py_INCREF(*outp);
-#else
-  if (theano_prep_output(outp, inp->ga.nd, inp->ga.dimensions, inp->ga.typecode, GA_C_ORDER, c) != 0)
+  if (params->inplace_output) {
+    Py_XDECREF(*outp);
+    *outp = inp;
+    Py_INCREF(*outp);
+  } else if (theano_prep_output(outp, inp->ga.nd, inp->ga.dimensions, inp->ga.typecode, GA_C_ORDER, c) != 0) {
     return 1;
-#endif
+  }
+
   if (theano_prep_output(x_mean, scale->ga.nd, scale->ga.dimensions, scale->ga.typecode, GA_C_ORDER, c) != 0)
     return 1;
   if (theano_prep_output(x_invstd, scale->ga.nd, scale->ga.dimensions, scale->ga.typecode, GA_C_ORDER, c) != 0)
@@ -43,30 +42,32 @@ int dnn_batchnorm_op(PyGpuArrayObject *inp, PyGpuArrayObject *scale,
   if (c_set_tensorNd(*outp, bn_output) != 0)
     return 1;
 
-#ifdef RUNNING_AVERAGES
-#ifdef INPLACE_RUNNING_MEAN
-  Py_XDECREF(*out_running_mean);
-  PyGpuArrayObject *running_mean = in_running_mean;
-  Py_INCREF(running_mean);
-#else
-  PyGpuArrayObject *running_mean = *out_running_mean;
-  running_mean = theano_try_copy(running_mean, in_running_mean);
-  if (running_mean == NULL) {
-    return 1;
-  }
-#endif
-#ifdef INPLACE_RUNNING_VAR
-  Py_XDECREF(*out_running_var);
-  PyGpuArrayObject *running_var = in_running_var;
-  Py_INCREF(running_var);
-#else
-  PyGpuArrayObject *running_var = *out_running_var;
-  running_var = theano_try_copy(running_var, in_running_var);
-  if (running_var == NULL) {
-    return 1;
+  PyGpuArrayObject *running_mean = NULL;
+  PyGpuArrayObject *running_var = NULL;
+  if (running_averages) {
+    if (params->inplace_running_mean) {
+      Py_XDECREF(*out_running_mean);
+      running_mean = in_running_mean;
+      Py_INCREF(running_mean);
+    } else {
+      running_mean = *out_running_mean;
+      running_mean = theano_try_copy(running_mean, in_running_mean);
+      if (running_mean == NULL) {
+        return 1;
+      }
+    }
+    if (params->inplace_running_var) {
+      Py_XDECREF(*out_running_var);
+      running_var = in_running_var;
+      Py_INCREF(running_var);
+    } else {
+      running_var = *out_running_var;
+      running_var = theano_try_copy(running_var, in_running_var);
+      if (running_var == NULL) {
+        return 1;
+      }
+    }
   }
-#endif
-#endif
 
   {
     const float falpha = 1.;
@@ -83,8 +84,8 @@ int dnn_batchnorm_op(PyGpuArrayObject *inp, PyGpuArrayObject *scale,
       beta = (void *)&fbeta;
     }
     cudnnStatus_t err = cudnnBatchNormalizationForwardTraining(
-      _handle,
-      MODE,
+      params->handle,
+      params->mode,
       alpha,
       beta,
       bn_input,
@@ -94,15 +95,9 @@ int dnn_batchnorm_op(PyGpuArrayObject *inp, PyGpuArrayObject *scale,
       bn_params,
       PyGpuArray_DEV_DATA(scale),
       PyGpuArray_DEV_DATA(bias),
-#ifdef RUNNING_AVERAGES
-      running_average_factor,
-      PyGpuArray_DEV_DATA(running_mean),
-      PyGpuArray_DEV_DATA(running_var),
-#else
-      0,
-      NULL,  // running mean, deliberately unused
-      NULL,  // running var, deliberately unused
-#endif
+      running_averages ? running_average_factor : 0,
+      running_averages ? PyGpuArray_DEV_DATA(running_mean) : NULL,
+      running_averages ? PyGpuArray_DEV_DATA(running_var): NULL,
       epsilon,
       PyGpuArray_DEV_DATA(*x_mean),
       PyGpuArray_DEV_DATA(*x_invstd)
@@ -112,10 +107,10 @@ int dnn_batchnorm_op(PyGpuArrayObject *inp, PyGpuArrayObject *scale,
                    cudnnGetErrorString(err));
       return 1;
     }
-#ifdef RUNNING_AVERAGES
-    *out_running_mean = running_mean;
-    *out_running_var = running_var;
-#endif
+    if (running_averages) {
+      *out_running_mean = running_mean;
+      *out_running_var = running_var;
+    }
   }
   return 0;
 }