[3DfftConv] Add stride value checking

theano/sandbox/cuda/opt.py

@@ -1121,7 +1121,6 @@ def local_gpu_softmax_with_bias(node):
 #### Convolution, maxpooling
 from theano.tensor.nnet import conv
 
-
 @register_opt()
 @local_optimizer([gpu_from_host, conv.ConvOp])
 def local_gpu_conv(node):
@@ -1261,7 +1260,6 @@ gpu_optimizer.register("conv_fft_full", local_conv_fft_full)
 
 
 from theano.sandbox.cuda.GpuConv3D import GpuConv3D
-
 def _gpu_conv3d_to_fftconv(node):
     # we import conv3d_fft locally to avoid pycuda warnings
     from theano.sandbox.cuda.fftconv import conv3d_fft
@@ -1272,31 +1270,29 @@ def _gpu_conv3d_to_fftconv(node):
     # Shuffle filters from (oc, 0, 1, t, ic) to (oc, ic, 0, 1, t)
     f = node.inputs[1]
     f = gpu_from_host(f.dimshuffle(0, 4, 1, 2, 3))
-
     rval = conv3d_fft(x, f)
-    
     # Shuffle from (oc, c, 0, 1, t) to (oc, 0, 1, t, c)
-    rval = gpu_from_host(rval.dimshuffle(0, 2, 3, 4, 1))
-
-
-    # Add biais
-    #rval = rval + node.inputs[2]
+    rval = gpu_from_host(rval.dimshuffle(0, 2, 3, 4, 1) + node.inputs[2])
 
     return rval
 
 @local_optimizer([GpuConv3D])
 def local_conv3d_fft(node):
 
-    if (isinstance(node.op, GpuConv3D)# and
-        #        node.inputs[3] == (1, 1, 1)]):
-        ):
+    try:
+        stride_x = tensor.get_scalar_constant_value(node.inputs[3][0])
+        stride_y = tensor.get_scalar_constant_value(node.inputs[3][1])
+        stride_z = tensor.get_scalar_constant_value(node.inputs[3][2])
+    except tensor.NotScalarConstantError:
+        pass
+    if (isinstance(node.op, GpuConv3D) and
+        (stride_x, stride_y, stride_z) == (1, 1, 1)):
         return [_gpu_conv3d_to_fftconv(node)]
 
 gpu_optimizer.register("conv3d_fft", local_conv3d_fft)
 
 
 from theano.sandbox.cuda.GpuConvGrad3D import GpuConvGrad3D
-
 def _gpu_convgrad3d_to_fftconv(node):
     # we import conv3d_fft locally to avoid pycuda warnings
     from theano.sandbox.cuda.fftconv import conv3d_fft
@@ -1307,7 +1303,6 @@ def _gpu_convgrad3d_to_fftconv(node):
     # Shuffle dCdH from (b, 0, 1, t, oc) to (oc, b, 0, 1, t)
     f = node.inputs[3]
     f = f.dimshuffle(4, 0, 1, 2, 3)
-
     rval = conv3d_fft(x, f)
     # Shuffle from (ic, oc, 0, 1, t) to (oc, 0, 1, t, ic)
     rval = gpu_from_host(rval.dimshuffle(1, 2, 3, 4, 0))
@@ -1317,15 +1312,20 @@ def _gpu_convgrad3d_to_fftconv(node):
 @local_optimizer([GpuConvGrad3D])
 def local_convgrad3d_fft(node):
 
-    if (isinstance(node.op, GpuConvGrad3D)# and
-        #        node.inputs[3] == (1, 1, 1)]):
-        ):
+    try:
+        stride_x = tensor.get_scalar_constant_value(node.inputs[3][0])
+        stride_y = tensor.get_scalar_constant_value(node.inputs[3][1])
+        stride_z = tensor.get_scalar_constant_value(node.inputs[3][2])
+    except tensor.NotScalarConstantError:
+        pass
+    if (isinstance(node.op, GpuConvGrad3D) and
+        (stride_x, stride_y, stride_z) == (1, 1, 1)):
         return [_gpu_convgrad3d_to_fftconv(node)]
+
 gpu_optimizer.register("convgrad3d_fft", local_convgrad3d_fft)
 
 
 from theano.sandbox.cuda.GpuConvTransp3D import GpuConvTransp3D
-
 def _gpu_convtransp3d_to_fftconv(node):
     # we import conv3d_fft locally to avoid pycuda warnings
     from theano.sandbox.cuda.fftconv import conv3d_fft
@@ -1336,8 +1336,7 @@ def _gpu_convtransp3d_to_fftconv(node):
     # Shuffle dCdH from (b, 0, 1, t, oc) to (b, oc, 0, 1, t)
     f = node.inputs[3]
     f = f.dimshuffle(0, 4, 1, 2, 3)
-
-    rval = conv3d_fft(f, x, border_mode = 'full')
+    rval = conv3d_fft(f, x, border_mode='full')
     # Shuffle from (ic, b, 0, 1, t) to (b, 0, 1, t, ic)
     rval = gpu_from_host(rval.dimshuffle(0, 2, 3, 4, 1))
 
@@ -1346,9 +1345,14 @@ def _gpu_convtransp3d_to_fftconv(node):
 @local_optimizer([GpuConvTransp3D])
 def local_convtransp3d_fft(node):
 
-    if (isinstance(node.op, GpuConvTransp3D)# and
-        #        node.inputs[3] == (1, 1, 1)]):
-        ):
+    try:
+        stride_x = tensor.get_scalar_constant_value(node.inputs[3][0])
+        stride_y = tensor.get_scalar_constant_value(node.inputs[3][1])
+        stride_z = tensor.get_scalar_constant_value(node.inputs[3][2])
+    except tensor.NotScalarConstantError:
+        pass
+    if (isinstance(node.op, GpuConvTransp3D) and
+        (stride_x, stride_y, stride_z) == (1, 1, 1)):
         return [_gpu_convtransp3d_to_fftconv(node)]
 
 gpu_optimizer.register("convtransp3d_fft", local_convtransp3d_fft)