diff --git a/onnx_tf/handlers/backend/depth_to_space.py b/onnx_tf/handlers/backend/depth_to_space.py index 5e3043a88..d64948984 100644 --- a/onnx_tf/handlers/backend/depth_to_space.py +++ b/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,15 +25,17 @@ 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") @@ -40,7 +43,9 @@ def _common(cls, node, **kwargs): 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): diff --git a/test/backend/test_dynamic_shape.py b/test/backend/test_dynamic_shape.py index 9b60c91f8..9e7384721 100644 --- a/test/backend/test_dynamic_shape.py +++ b/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(