Backport dynamic shape support for ConvTranspose and BatchNormalization (#724)

winnietsang · web-flow · commit c6784d414e6f · 2020-08-20T15:18:44.000-07:00
1. Backport dynamic shape support for ConvTranspose and
BatchNormalization to tf-1.x branch
2. Backport convTranspose testcases to tf-1.x branch
diff --git a/onnx_tf/handlers/backend/batch_normalization.py b/onnx_tf/handlers/backend/batch_normalization.py
@@ -29,6 +29,10 @@ def _common(cls, node, **kwargs):
 
     params_shape_broadcast = list([1, x_shape[1]] +
                                   [1 for _ in range(2, x_rank)])
+    # process unknown channel shape
+    if params_shape_broadcast[1] is None:
+      params_shape_broadcast[1] = tf.shape(x)[1]
+      params_shape_broadcast = tf.stack(params_shape_broadcast)
 
     total_num_dim = len(x.get_shape())
     scale = tf.reshape(tensor_dict[node.inputs[1]], params_shape_broadcast)
diff --git a/onnx_tf/handlers/backend/conv_mixin.py b/onnx_tf/handlers/backend/conv_mixin.py
@@ -127,6 +127,16 @@ def conv(cls, node, input_dict, transpose=False):
             ]
           conv_output_shape.insert(compute_c_idx, weights_shape[-2])
 
+          def handle_dynamic_batch_size(output_shape, batch_idx):
+            output_shape[batch_idx] = tf.shape(x)[batch_idx]
+            return tf.stack(output_shape)
+
+          # process dynamic batch size
+          if conv_output_shape[storage_format.find("N")] is None:
+            batch_idx = storage_format.find("N")
+            conv_output_shape = handle_dynamic_batch_size(conv_output_shape,
+                    batch_idx)
+
           # make strides to match input rank
           strides_full = [1] + strides
           strides_full.insert(compute_c_idx, 1)
@@ -169,6 +179,12 @@ def conv(cls, node, input_dict, transpose=False):
               pads[spatial_format.find(d) + spatial_size]
               for d, s in zip(compute_format, conv_rs_shape)
           ]
+
+          # process dynamic batch size
+          if size[compute_format.find("N")] is None:
+            batch_idx = compute_format.find("N")
+            size = handle_dynamic_batch_size(size, batch_idx)
+
           conv_rs = tf.slice(conv_rs, begin=begin, size=size)
 
           convolved.append(conv_rs)
@@ -190,6 +206,12 @@ def conv(cls, node, input_dict, transpose=False):
             ]
           conv_output_shape.insert(compute_c_idx, weights_shape[-2])
 
+          # process dynamic batch size
+          if conv_output_shape[storage_format.find("N")] is None:
+            batch_idx = storage_format.find("N")
+            conv_output_shape = handle_dynamic_batch_size(conv_output_shape,
+                    batch_idx)
+
           # make strides to match input rank
           strides_full = [1] + strides
           strides_full.insert(compute_c_idx, 1)
diff --git a/onnx_tf/handlers/backend/upsample.py b/onnx_tf/handlers/backend/upsample.py
@@ -9,6 +9,7 @@
 from onnx_tf.handlers.handler import partial_support
 from onnx_tf.handlers.handler import ps_description
 from onnx_tf.handlers.handler import tf_func
+from onnx_tf.common.tf_helper import tf_shape
 
 
 @onnx_op("Upsample")
@@ -54,7 +55,7 @@ def version_7(cls, node, **kwargs):
   @classmethod
   def version_9(cls, node, **kwargs):
     x = kwargs["tensor_dict"][node.inputs[0]]
-    x_shape = x.get_shape().as_list()
+    x_shape = tf_shape(x)
     attrs = copy.deepcopy(node.attrs)
     scales = kwargs["tensor_dict"][node.inputs[1]]
 
@@ -65,7 +66,8 @@ def version_9(cls, node, **kwargs):
     with tf.control_dependencies([assert_n_c_scale_is_one]):
       h_w_scale = scales[2:]
       h_w_shape = x_shape[2:]
-      new_h_w_shape = tf.cast(h_w_scale * h_w_shape, tf.int32)
+      new_h_w_shape = tf.cast(h_w_scale * tf.cast(h_w_shape, scales.dtype),
+                              tf.int32)
 
       mode = attrs.get("mode", "nearest")
       if mode.lower() == "bilinear" or mode.lower() == "linear":
diff --git a/test/backend/test_dynamic_shape.py b/test/backend/test_dynamic_shape.py
@@ -23,7 +23,7 @@ class TestDynamicShape(unittest.TestCase):
 
   def _get_rnd_float32(self, low=-1.0, high=1.0, shape=None):
     output = np.random.uniform(low, high, shape)
-    if shape == None:
+    if shape is None:
       return np.float32(output)
     else:
       return output.astype(np.float32)
@@ -87,6 +87,100 @@ def test_arg_min(self):
     expected_output = x.shape[axis] - expected_output - 1
     np.testing.assert_almost_equal(output['Y'], expected_output)
 
+  def _batch_normalization(self, x, mean, variance, bias, scale,
+                           variance_epsilon):
+    inv = np.reciprocal(np.sqrt(variance + variance_epsilon))
+    if scale is not None:
+      inv *= scale
+    return x * inv + (bias - mean * inv if bias is not None else -mean * inv)
+
+  def test_batch_normalization(self):
+    if legacy_opset_pre_ver(6):
+      raise unittest.SkipTest("Backend doesn't support consumed flag")
+    node_def = helper.make_node("BatchNormalization",
+                                ["X", "scale", "bias", "mean", "var"], ["Y"],
+                                epsilon=0.001)
+    graph_def = helper.make_graph(
+        [node_def],
+        name="test_unknown_shape",
+        inputs=[
+            helper.make_tensor_value_info("X", TensorProto.FLOAT, [None, None, None, None]),
+            helper.make_tensor_value_info("scale", TensorProto.FLOAT, [None]),
+            helper.make_tensor_value_info("bias", TensorProto.FLOAT, [None]),
+            helper.make_tensor_value_info("mean", TensorProto.FLOAT, [None]),
+            helper.make_tensor_value_info("var", TensorProto.FLOAT, [None])
+        ],
+        outputs=[
+            helper.make_tensor_value_info("Y", TensorProto.FLOAT, [None, None, None, None])
+        ])
+    x_shape = [3, 5, 4, 2]
+    param_shape = [5]
+    _param_shape = [1, 5, 1, 1]
+    x = self._get_rnd_float32(0, 1, shape=x_shape)
+    m = self._get_rnd_float32(0, 1, shape=param_shape)
+    _m = m.reshape(_param_shape)
+    v = self._get_rnd_float32(0, 1, shape=param_shape)
+    _v = v.reshape(_param_shape)
+    scale = self._get_rnd_float32(0, 1, shape=param_shape)
+    _scale = scale.reshape(_param_shape)
+    bias = self._get_rnd_float32(0, 1, shape=param_shape)
+    _bias = bias.reshape(_param_shape)
+    golden = self._batch_normalization(x, _m, _v, _bias, _scale, 0.001)
+    tf_rep = onnx_graph_to_tensorflow_rep(graph_def)
+    output = tf_rep.run({"X": x, "scale": scale, "bias": bias, "mean": m, "var": v})
+    np.testing.assert_almost_equal(output["Y"], golden, decimal=5)  
+
+  def test_conv_transpose(self):
+    # test dynamic batch size on transpose of 2d convolution
+    pads = [1, 1, 1, 1]
+    x_shape = [1, 3, 4, 6]
+    x = self._get_rnd_float32(shape=x_shape)
+    weight_shape = [3, 5, 2, 2]
+    weights = self._get_rnd_float32(shape=weight_shape)
+
+    node_def = helper.make_node("ConvTranspose", ["X", "weights"], ["Y"],
+                                pads=pads)
+    graph_def = helper.make_graph(
+        [node_def],
+        name="test_unknown_shape",
+        inputs=[
+            helper.make_tensor_value_info("X", TensorProto.FLOAT, [None, 3, 4, 6]),
+            helper.make_tensor_value_info("weights", TensorProto.FLOAT, weight_shape)
+        ],
+        outputs=[
+            helper.make_tensor_value_info("Y", TensorProto.FLOAT, [None, None, None, None])
+        ])
+
+    tf_rep = onnx_graph_to_tensorflow_rep(graph_def)
+    output = tf_rep.run({"X": x, "weights": weights})
+
+    padh_left = weight_shape[2] - 1 - pads[0]
+    padh_right = weight_shape[2] - 1 - pads[1]
+    padw_left = weight_shape[3] - 1 - pads[2]
+    padw_right = weight_shape[3] - 1 - pads[3]
+
+    kh = weight_shape[2]
+    kw = weight_shape[3]
+    outh = x_shape[2] + padh_right + padh_right - (kh - 1)
+    outw = x_shape[3] + padw_right + padw_right - (kw - 1)
+
+    out_shape = [x_shape[0], weight_shape[1], outh, outw]
+
+    test_output = np.zeros(out_shape)
+    for b in range(0, x_shape[0]):
+      for m in range(0, weight_shape[1]):
+        for c in range(0, x_shape[1]):
+          for h in range(0, outh):
+            for w in range(0, outw):
+              for k1 in range(h, h + kh):
+                for k2 in range(w, w + kw):
+                  if (k1 - padh_left >= 0 and k2 - padw_left >= 0):
+                    test_output[b][m][h][w] += x[b][c][k1 - padh_left][
+                        k2 - padw_left] * weights[c][m][kh + h - 1 -
+                                                        k1][kw + w - 1 - k2]
+
+    np.testing.assert_almost_equal(output["Y"], test_output, decimal=5)
+
   def test_slice(self):
     # test case 1 with normal inputs
     axes = [0, 1, 2]
diff --git a/test/backend/test_node.py b/test/backend/test_node.py
@@ -23,7 +23,7 @@ class TestNode(unittest.TestCase):
 
   def _get_rnd_float32(self, low=-1.0, high=1.0, shape=None):
     output = np.random.uniform(low, high, shape)
-    if shape == None:
+    if shape is None:
       return np.float32(output)
     else:
       return output.astype(np.float32)
@@ -460,29 +460,68 @@ def test_conv_integer(self):
     np.testing.assert_almost_equal(output["Y"], y)
 
   def test_conv_transpose(self):
-    # Fix test in the future.
-    return
-    device = "CUDA"
-    if not supports_device(device):
-      raise unittest.SkipTest(
-          "Backend doesn't support device {}".format(device))
+    device = "CUDA" if supports_device("CUDA") else "CPU"
+
+    pads = [1, 1]
     node_def = helper.make_node("ConvTranspose", ["X", "weights"], ["Y"],
-                                pads=[1, 1])
-    x_shape = [1, 5, 4]
+                                pads=pads)
+    x_shape = [1, 3, 4]
     x = self._get_rnd_float32(shape=x_shape)
-    weight_shape = [5, 3, 2]
+    weight_shape = [3, 5, 2]
     weights = self._get_rnd_float32(shape=weight_shape)
     output = run_node(node_def, [x, weights], device=device)
-    out_shape = [x_shape[0], weight_shape[1], x_shape[2]]
+
+    padh_left =  weight_shape[2]-1-pads[0]
+    padh_right = weight_shape[2]-1-pads[1]
+    kh = weight_shape[2]
+    outh = x_shape[2] + padh_right + padh_right - (kh - 1)
+
+    out_shape = [x_shape[0], weight_shape[1], outh]
+
     test_output = np.zeros(out_shape)
     for b in range(0, x_shape[0]):
       for m in range(0, weight_shape[1]):
-        for h in range(0, x_shape[2]):
-          v = 0
-          for c in range(0, x_shape[1]):
-            for k in range(h, min(h + weight_shape[2], x_shape[2])):
-              v += x[b][c][k] * weights[c][m][k - h]
-          test_output[b][m][h] = v
+        for c in range(0, x_shape[1]):
+          for h in range(0, outh):
+            for k in range(h , h + kh):
+              if (k - padh_left >= 0):
+                test_output[b][m][h] += x[b][c][k-padh_left] * weights[c][m][kh+h-1-k]
+
+    np.testing.assert_almost_equal(output["Y"], test_output, decimal=5)
+
+    # test for spatial dimension of colnolution is 2
+    pads = [1, 1, 1, 1]
+    node_def = helper.make_node("ConvTranspose", ["X", "weights"], ["Y"],
+                                pads=pads)
+    x_shape = [1, 3, 4, 6]
+    x = self._get_rnd_float32(shape=x_shape)
+    weight_shape = [3, 5, 2, 2]
+    weights = self._get_rnd_float32(shape=weight_shape)
+    output = run_node(node_def, [x, weights],device=device)
+
+    padh_left =  weight_shape[2]-1-pads[0]
+    padh_right = weight_shape[2]-1-pads[1]
+    padw_left =  weight_shape[3]-1-pads[2]
+    padw_right = weight_shape[3]-1-pads[3]
+
+    kh = weight_shape[2]
+    kw = weight_shape[3]
+    outh = x_shape[2] + padh_right + padh_right - (kh - 1)
+    outw = x_shape[3] + padw_right + padw_right - (kw - 1)
+
+    out_shape = [x_shape[0], weight_shape[1], outh, outw]
+
+    test_output = np.zeros(out_shape)
+    for b in range(0, x_shape[0]):
+      for m in range(0, weight_shape[1]):
+        for c in range(0, x_shape[1]):
+          for h in range(0, outh):
+            for w in range(0, outw):
+              for k1 in range(h , h + kh):
+                for k2 in range(w , w + kw):
+                  if (k1 - padh_left >= 0 and k2 - padw_left >= 0):
+                    test_output[b][m][h][w] += x[b][c][k1-padh_left][k2-padw_left] * weights[c][m][kh+h-1-k1][kw+w-1-k2]
+
     np.testing.assert_almost_equal(output["Y"], test_output, decimal=5)
 
   def test_cosh(self):
@@ -1124,7 +1163,7 @@ def test_loop(self):
     x_in = helper.make_tensor_value_info('x', TensorProto.INT32, [None])
     y_in = helper.make_tensor_value_info('y', TensorProto.INT32, [None])
 
-    cond_out = helper.make_tensor_value_info('cond', TensorProto.STRING, [])
+    cond_out = helper.make_tensor_value_info('cond', TensorProto.BOOL, [])
     new_cond_out = helper.make_tensor_value_info('new_cond', TensorProto.BOOL,
                                                  [])
     sum1_out = helper.make_tensor_value_info('sum1', TensorProto.INT32, [None])
