Improve ConvTranspose shape inference to work with unknown batch dim

Signed-off-by: Tom Wildenhain <[email protected]>

Author: TomWildenhain-Microsoft

tests/test_backend.py

@@ -3257,6 +3257,31 @@ def func(filter_val, out_backprop_val):
         out_backprop_val = np.random.randint(low=0, high=256, size=[1, 5, 5, 5]).astype(np.float32)
         self._run_test_case(func, [_OUTPUT], {_INPUT: filters_val, _INPUT1: out_backprop_val})
 
+    @check_tf_min_version("1.15", "tf.repeat needs tf 1.15")
+    @check_opset_min_version(10, "Conv2DBackpropInput")
+    def test_Conv2DBackpropInput_shape_implied(self):
+        batch_dim_val = np.array(1, dtype=np.int32)
+        def func(filter_val, out_backprop_val, batch_dim):
+            out_backprop_val = tf.repeat(out_backprop_val, batch_dim, axis=0)
+            s = tf.shape(out_backprop_val)
+            t1 = tf.constant([0], dtype=tf.int32)
+            t2 = tf.constant([1], dtype=tf.int32)
+            batch_dim = tf.strided_slice(s, t1, t2, shrink_axis_mask=1)
+            # Sometimes the size given is a stack of constants with unknown batch dim
+            input_sizes_val = tf.stack([batch_dim, 10, 10, 3])
+            return conv2d_backprop_input(input_sizes=input_sizes_val, filter=filter_val,
+                                         out_backprop=out_backprop_val, strides=[1, 2, 2, 1],
+                                         padding='SAME', name=_TFOUTPUT)
+        filters_val = np.random.randint(low=0, high=256, size=[3, 3, 3, 5]).astype(np.float32)
+        out_backprop_val = np.random.randint(low=0, high=256, size=[1, 5, 5, 5]).astype(np.float32)
+        def graph_validator(g):
+            for n in g.get_nodes():
+                if n.type == 'ConvTranspose':
+                    return "output_shape" in n.attr
+            return False
+        self._run_test_case(func, [_OUTPUT], {_INPUT: filters_val, _INPUT1: out_backprop_val, _INPUT2: batch_dim_val},
+                            graph_validator=graph_validator)
+
     @check_opset_min_version(10, "Conv2DBackpropInput")
     def test_Conv2DBackpropInput_const_valid(self):
         input_sizes_val_ = np.array([1, 12, 12, 3], dtype=np.int32)

tf2onnx/onnx_opset/nn.py

@@ -366,6 +366,27 @@ def version_11(cls, ctx, node, **kwargs):
         # No change.
         cls.version_1(ctx, node, **kwargs)
 
+def get_shape_from_const_or_concat(ctx, node):
+    if node.is_const():
+        return node.get_tensor_value()
+    if node.type == 'Concat':
+        # Sometimes the shape is formed by concating a bunch of consts together
+        res = []
+        if any(ctx.get_shape(inp) != [1] for inp in node.input):
+            return None
+        for i, inp in enumerate(node.inputs):
+            # The concat is converted from a Pack. Conversion adds an unsqueeze to the inputs.
+            if node.inputs[i].type == 'Unsqueeze' and node.inputs[i].inputs[0].is_scalar():
+                res.append(node.inputs[i].inputs[0].get_tensor_value())
+            else:
+                if i == 0:
+                    # For the batch dimension we don't care if it is unknown
+                    res.append(-1)
+                else:
+                    return None
+        return res
+    return None
+
 @tf_op(["Conv2DBackpropInput", "Conv3DBackpropInputV2"])
 class ConvTranspose:
     @classmethod
@@ -386,8 +407,9 @@ def version_1(cls, ctx, node, **kwargs):
         output_shape_orig = node.output_shapes
 
         # ouput_shape is explicitly specified here, in this case pads values are auto generated/calculated.
-        if node.inputs[0].is_const():
-            output_shape = ctx.get_shape(node.output[0])
+        output_shape = get_shape_from_const_or_concat(ctx, node.inputs[0])
+        if output_shape is not None:
+            #output_shape = ctx.get_shape(node.output[0])
             if is_channels_last(node):
                 new_output_shape = [output_shape[1], output_shape[2]]
                 input_dims = [input_shape[1], input_shape[2]]
@@ -407,6 +429,7 @@ def version_1(cls, ctx, node, **kwargs):
 
             node.set_attr("output_shape", new_output_shape)
         else:
+            # FIXME: This case fails in edge cases where strides > 1
             input_shape = ctx.make_node("Cast", [node.input[0]], attr={'to': TensorProto.INT64})
             output_shape = ctx.make_node("Shape", [node.output[0]])
             output_h = GraphBuilder(ctx).make_slice(

tf2onnx/onnx_opset/tensor.py

@@ -1054,7 +1054,8 @@ def version_1(cls, ctx, node, **kwargs):
         # insert Unsqueeze on each input
         for i, n in enumerate(node.inputs):
             dtype = ctx.get_dtype(node.input[i])
-            shape = ctx.get_shape(node.input[i])
+            shape = ctx.get_shape(node.input[i]).copy()
+            shape.insert(axis, 1)
             new_node = ctx.make_node("Unsqueeze", [node.input[i]], op_name_scope=node.name, attr={"axes": [axis]},
                                      shapes=[shape], dtypes=[dtype])
             output_name = new_node.output[0]