Separate spatialtf gradients, for inputs and transformation, into separate Ops

theano/gpuarray/c_code/dnn_sptf_gi.c

+#section support_code_struct
+
+int
+dnn_sptf_gi(PyGpuArrayObject * input,
+            PyGpuArrayObject * theta,
+            PyGpuArrayObject * grid,
+            PyArrayObject * grid_dims,
+            PyGpuArrayObject * dy,
+            cudnnSpatialTransformerDescriptor_t desc,
+            double alpha, double beta,
+            PyGpuArrayObject ** input_grad,
+            PyGpuArrayObject ** grid_grad,
+            cudnnHandle_t _handle)
+{
+    PyErr_SetString(PyExc_NotImplementedError, "Gradient for spatial transformer is not yet implemented.");
+    return -1;
+}

theano/gpuarray/c_code/dnn_sptf_grad.c

-#section support_code_struct
-
-int
-dnn_sptf_grad(PyGpuArrayObject * input,
-              PyGpuArrayObject * theta,
-              PyGpuArrayObject * grid,
-              PyArrayObject * grid_dims,
-              PyGpuArrayObject * dy,
-              cudnnSpatialTransformerDescriptor_t desc,
-              double alpha, double beta,
-              PyGpuArrayObject ** output_grad,
-              PyGpuArrayObject ** grid_grad,
-              cudnnHandle_t _handle)
-{
-    PyErr_SetString(PyExc_NotImplementedError, "Gradient for spatial transformer is not yet implemented.");
-    return -1;
-}

theano/gpuarray/c_code/dnn_sptf_gt.c

+#section support_code_struct
+
+int
+dnn_sptf_gt(PyGpuArrayObject * dgrid,
+            cudnnSpatialTransformerDescriptor_t desc,
+            PyGpuArrayObject ** dtheta,
+            cudnnHandle_t _handle)
+{
+    PyErr_SetString(PyExc_NotImplementedError, "Gradient for spatial transformer is not yet implemented.");
+    return -1;
+}

theano/gpuarray/dnn.py

@@ -2937,38 +2937,67 @@ class GpuDnnTransformer(DnnBase):
         _, grid = outputs
         dy = grads[0]
 
-        gop = GpuDnnTransformerGrad(self.dtype)(img, theta, grid, grid_dims, dy,
+        dimg, dgrid = GpuDnnTransformerGradI(self.dtype)(img, theta, grid, grid_dims, dy,
                                                          desc, alpha, beta)
+        dtheta = GpuDnnTransformerGradT(self.dtype)(dgrid, desc)
 
-        return gop
+        return [dimg, dtheta,
+                theano.gradient.grad_undefined(self, 2, grid_dims),
+                theano.gradient.grad_undefined(self, 3, desc),
+                theano.gradient.grad_undefined(self, 4, alpha),
+                theano.gradient.grad_undefined(self, 5, beta)]
 
 
-class GpuDnnTransformerGrad(DnnBase):
+class GpuDnnTransformerGradI(DnnBase):
+    """
+    Gradients of inputs of the spatial transformer
+    """
     __props__ = ('dtype',)
     _cop_num_inputs = 8
-    _cop_num_outputs = 2
+    _cop_num_outputs = 1
     _f16_ok = True
 
     def __init__(self, dtype=theano.config.floatX):
-        DnnBase.__init__(self, ["c_code/dnn_sptf_grad.c"], "dnn_sptf_grad")
+        DnnBase.__init__(self, ["c_code/dnn_sptf_gi.c"], "dnn_sptf_gi")
         self.dtype = dtype
 
     def make_node(self, img, theta, grid, grid_dims, dy, desc, alpha, beta):
         context_name = infer_context_name(img)
 
+        dy = as_gpuarray_variable(dy, context_name)
         dimg = GpuArrayType(dtype=self.dtype,
                             broadcastable=img.type.ndim * (False,),
                             context_name=context_name)()
-        dtheta = GpuArrayType(dtype=self.dtype,
-                              broadcastable=theta.type.ndim * (False,),
+        dgrid = GpuArrayType(dtype=self.dtype,
+                             broadcastable=img.type.ndim * (False,),
                              context_name=context_name)()
 
         inputs = [img, theta, grid, grid_dims, dy, desc, alpha, beta]
-        outputs = [dimg, dtheta,
-                   theano.gradient.grad_undefined(self, 2, grid_dims),
-                   theano.gradient.grad_undefined(self, 3, desc),
-                   theano.gradient.grad_undefined(self, 4, alpha),
-                   theano.gradient.grad_undefined(self, 5, beta)]
+        outputs = [dimg, dgrid]
+
+        return Apply(self, inputs, outputs)
+
+
+class GpuDnnTransformerGradT(DnnBase):
+    """
+    Gradients of the affine transformation generated by the localisation network
+    """
+    __props__ = ('dtype',)
+    _cop_num_inputs = 2
+    _cop_num_outputs = 1
+    _f16_ok = True
+
+    def __init__(self, dtype=theano.config.floatX):
+        DnnBase.__init__(self, ["c_code/dnn_sptf_gt.c"], "dnn_sptf_gt")
+        self.dtype = dtype
+
+    def make_node(self, dgrid, desc):
+        context_name = infer_context_name(dgrid)
+        dtheta = GpuArrayType(dtype=self.dtype,
+                              broadcastable=(dgrid.type.ndim - 1) * (False,),
+                              context_name=context_name)()
+        inputs = [dgrid, desc]
+        outputs = [dtheta]
 
         return Apply(self, inputs, outputs)