Make SubPixelUpscaling work with any backend/dim ordering combination

keras_contrib/applications/densenet.py

@@ -651,11 +651,3 @@ def __create_fcn_dense_net(nb_classes, img_input, include_top, nb_dense_block=5,
         x = Reshape((row, col, nb_classes))(x)
 
     return x
-
-if __name__ == '__main__':
-    model = DenseNetFCN((32, 32, 3), growth_rate=16, nb_layers_per_block=[4, 5, 7, 10, 12, 15],
-                        dropout_rate=0.2, upsampling_type='subpixel')
-
-    from keras.utils.visualize_util import plot
-
-    plot(model, to_file='densenet fcn.png', show_shapes=True)

keras_contrib/backend/tensorflow_backend.py

@@ -17,6 +17,9 @@ from keras.backend.tensorflow_backend import _preprocess_conv3d_kernel
 from keras.backend.tensorflow_backend import _preprocess_border_mode
 from keras.backend.tensorflow_backend import _postprocess_conv3d_output
 from keras.backend.tensorflow_backend import _preprocess_border_mode
+from keras.backend.tensorflow_backend import _preprocess_conv2d_input
+from keras.backend.tensorflow_backend import _postprocess_conv2d_output
+
 py_all = all
 
 
@@ -107,10 +110,11 @@ def extract_image_patches(x, ksizes, ssizes, border_mode="same",
 
 def depth_to_space(input, scale):
     ''' Uses phase shift algorithm to convert channels/depth for spatial resolution '''
-    assert K.image_dim_ordering() == 'tf', 'depth_to_space backend function can only be used with "tf" dim ' \
-                                           'ordering when using tensorflow backend'
 
-    return tf.depth_to_space(input, scale)
+    input = _preprocess_conv2d_input(input, image_dim_ordering())
+    out = tf.depth_to_space(input, scale)
+    out = _postprocess_conv2d_output(out, image_dim_ordering())
+    return out
 
 
 def moments(x, axes, shift=None, keep_dims=False):

keras_contrib/backend/theano_backend.py

@@ -24,6 +24,8 @@ from keras.backend.theano_backend import _preprocess_conv3d_kernel
 from keras.backend.theano_backend import _preprocess_conv3d_filter_shape
 from keras.backend.theano_backend import _preprocess_border_mode
 from keras.backend.theano_backend import _postprocess_conv3d_output
+from keras.backend.theano_backend import _preprocess_conv2d_input
+from keras.backend.theano_backend import _postprocess_conv2d_output
 
 import itertools
 
@@ -121,11 +123,11 @@ def extract_image_patches(X, ksizes, strides, border_mode="valid", dim_ordering=
 
 def depth_to_space(input, scale):
     ''' Uses phase shift algorithm to convert channels/depth for spatial resolution '''
-    assert K.image_dim_ordering() == 'th', 'depth_to_space backend function can only be used with "th" dim ' \
-                                           'ordering when using theano backend'
+
+    input = _preprocess_conv2d_input(input, image_dim_ordering())
 
     b, k, row, col = input.shape
-    output_shape = (b, input._keras_shape[1] // (scale ** 2), row * scale, col * scale)
+    output_shape = (b, k // (scale ** 2), row * scale, col * scale)
 
     out = T.zeros(output_shape)
     r = scale
@@ -133,6 +135,7 @@ def depth_to_space(input, scale):
     for y, x in itertools.product(range(scale), repeat=2):
         out = T.inc_subtensor(out[:, :, y::r, x::r], input[:, r * y + x:: r * r, :, :])
 
+    out = _postprocess_conv2d_output(out, input, None, None, None, image_dim_ordering())
     return out
 
 

keras_contrib/layers/convolutional.py

@@ -479,10 +479,6 @@ class SubPixelUpscaling(Layer):
     This layer performs the depth to space operation on the convolution filters, and returns a
     tensor with the size as defined below.
 
-    # Note:
-        This layer does not work with mismatched backend and dim ordering.
-        For example, TH dim ordering with Tensorflow backend / TF dim ordering with Theano backend.
-
     # Example :
     ```python
         # A standard subpixel upscaling block
@@ -502,8 +498,7 @@ class SubPixelUpscaling(Layer):
 
     # Arguments
         scale_factor: Upscaling factor.
-        dim_ordering: Can be 'th' or 'tf'. Note: mismatched dim ordering will
-                      cause a ValueError to be raised.
+        dim_ordering: Can be 'default', 'th' or 'tf'.
 
     # Input shape
         4D tensor with shape:
@@ -528,11 +523,6 @@ class SubPixelUpscaling(Layer):
         if self.dim_ordering == 'default':
             self.dim_ordering = K.image_dim_ordering()
 
-        if (K.backend() == 'theano' and self.dim_ordering == 'tf') or \
-                (K.backend() == 'tensorflow' and self.dim_ordering == 'th'):
-            raise ValueError('SubPixelUpscaling cannot be used with mismatched backend / dim ordering combinations. '
-                             'Backend : %s, Dim Ordering : %s' % (K.backend(), self.dim_ordering))
-
     def build(self, input_shape):
         pass
 

tests/keras_contrib/layers/test_convolutional.py

@@ -157,11 +157,6 @@ def test_sub_pixel_upscaling():
     for scale_factor in [2, 3, 4]:
         input_data = np.random.random((nb_samples, 4 * (scale_factor ** 2), nb_row, nb_col))
 
-        if K.backend() == 'theano':
-            K.set_image_dim_ordering('th')
-        else:
-            K.set_image_dim_ordering('tf')
-
         if K.image_dim_ordering() == 'tf':
             input_data = input_data.transpose((0, 2, 3, 1))