fix dynamic shape inference in DepthToSpace (#880)

* fix dynamic shape inference in DepthToSpace

Signed-off-by: Dom Miketa <[email protected]>

* add unittest

Signed-off-by: Dom Miketa <[email protected]>

* formatting with yapf and pylint

Signed-off-by: Dom Miketa <[email protected]>

* format unittest

Signed-off-by: Dom Miketa <[email protected]>

Author: d-miketa

onnx_tf/handlers/backend/depth_to_space.py

@@ -3,6 +3,7 @@
 import tensorflow as tf
 
 from onnx_tf.common import get_data_format
+from onnx_tf.common.tf_helper import tf_shape
 from onnx_tf.handlers.backend_handler import BackendHandler
 from onnx_tf.handlers.handler import onnx_op
 from onnx_tf.handlers.handler import tf_func
@@ -24,23 +25,27 @@ def version_1(cls, node, **kwargs):
     attrs = copy.deepcopy(node.attrs)
     attrs["data_format"] = storage_format
     return [
-        cls.make_tensor_from_onnx_node(
-            node, attrs=attrs, c_first_cuda_only=True, **kwargs)
+        cls.make_tensor_from_onnx_node(node,
+                                       attrs=attrs,
+                                       c_first_cuda_only=True,
+                                       **kwargs)
     ]
 
   @classmethod
   def _common(cls, node, **kwargs):
     x = kwargs["tensor_dict"][node.inputs[0]]
     x_rank = len(x.get_shape())
-    storage_format, compute_format = get_data_format(x_rank)
+    storage_format, _ = get_data_format(x_rank)
     attrs = copy.deepcopy(node.attrs)
     attrs["data_format"] = storage_format
     mode = attrs.get("mode", "DCR")
 
     if mode == "CRD":
       # need native computation
       bsize = attrs.get("blocksize")
-      batch, channel, height, width = x.shape
+      x_shape = tf_shape(x)
+      batch, channel = x_shape[0], x_shape[1]
+      height, width = x_shape[2], x_shape[3]
       csize = channel // (bsize**2)
 
       reshape_node = tf.reshape(x, [batch, csize, bsize, bsize, height, width])
@@ -50,11 +55,12 @@ def _common(cls, node, **kwargs):
                      [batch, csize, height * bsize, width * bsize])
       ]
 
-    else:
-      return [
-          cls.make_tensor_from_onnx_node(
-              node, attrs=attrs, c_first_cuda_only=True, **kwargs)
-      ]
+    return [
+        cls.make_tensor_from_onnx_node(node,
+                                       attrs=attrs,
+                                       c_first_cuda_only=True,
+                                       **kwargs)
+    ]
 
   @classmethod
   def version_11(cls, node, **kwargs):

test/backend/test_dynamic_shape.py

@@ -248,6 +248,37 @@ def test_conv_transpose(self):
 
     np.testing.assert_almost_equal(tf_model_output[0], test_output, decimal=5)
 
+  def test_depth_to_space(self):
+    b, c, h, w = shape = [2, 48, 5, 6]
+    blocksize = 4
+    x = self._get_rnd_float32(shape=shape)
+    node_def = helper.make_node("DepthToSpace", ["X"], ["Y"],
+                                blocksize=blocksize,
+                                mode="DCR")
+    graph_def = helper.make_graph(
+        [node_def],
+        name="test_unknown_shape",
+        inputs=[
+            helper.make_tensor_value_info("X", TensorProto.FLOAT,
+                                          [None, None, None, None])
+        ],
+        outputs=[
+            helper.make_tensor_value_info("Y", TensorProto.FLOAT,
+                                          [None, None, None, None])
+        ])
+    tf_rep = onnx_graph_to_tensorflow_rep(graph_def)
+    # export to tf.saved_model
+    model_path = 'test_dynamic_shape/depth_to_space'
+    tf_rep.export_graph(model_path)
+    # load the saved_model back
+    tf_model = tf.saved_model.load(model_path)
+    # run the model
+    tf_model_output = tf_model(X=x)
+    tmp = np.reshape(x, [b, blocksize, blocksize, c // (blocksize**2), h, w])
+    tmp = np.transpose(tmp, [0, 3, 4, 1, 5, 2])
+    y = np.reshape(tmp, [b, c // (blocksize**2), h * blocksize, w * blocksize])
+    np.testing.assert_almost_equal(tf_model_output[0], y)
+
   def test_eye_like(self):
     if legacy_opset_pre_ver(9):
       raise unittest.SkipTest("ONNX version {} doesn't support EyeLike.".format(