Minor bugfixes to support tflite (#1275)

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

Author: TomWildenhain-Microsoft

tf2onnx/onnx_opset/controlflow.py

@@ -268,7 +268,8 @@ def version_1(cls, ctx, node, **kwargs):
             branches[branch] = g
 
         _ = ctx.make_node("If", node.input[:1], name=node.name, output_count=len(output_shapes),
-                          shapes=output_shapes, dtypes=output_dtypes, skip_conversion=True, branches=branches)
+                          shapes=output_shapes, dtypes=output_dtypes, outputs=node.output, skip_conversion=True,
+                          branches=branches)
 
 
 @tf_op(["TensorListSetItem"])
@@ -629,6 +630,7 @@ def parameter_binding(g, inputs, state_vars=None):
         else:
             binding[k] = inputs[i]
             i += 1
+    utils.make_sure(i == len(inputs), "Parameter count mismatch while binding controlflow")
     return binding
 
 

tf2onnx/onnx_opset/nn.py

@@ -577,6 +577,10 @@ def version_1(cls, ctx, node, **kwargs):
         if len(kernel_shape) != 4:
             raise ValueError("only Conv2D is supported")
         k_h, k_w, k_input_channels, k_channel_multiplier = kernel_shape
+        if "depth_multiplier" in node.attr:
+            depth_multiplier = node.get_attr_int("depth_multiplier")
+            k_input_channels //= depth_multiplier
+            k_channel_multiplier *= depth_multiplier
         if k_input_channels < 1:
             raise ValueError("input channel must be positive")
         k_output_channels = k_input_channels * k_channel_multiplier

tf2onnx/onnx_opset/tensor.py

@@ -1148,8 +1148,10 @@ def any_version(cls, opset, 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]).copy()
-            shape.insert(axis, 1)
+            shape = ctx.get_shape(node.input[i])
+            if shape is not None:
+                shape = shape.copy()
+                shape.insert(axis, 1)
             new_node = gb.make_unsqueeze(
                 {'data': node.input[i], 'axes': [axis]},
                 op_name_scope=node.name, shapes=[shape], dtypes=[dtype], return_node=True)

tf2onnx/rewriter/eye_rewriter.py

@@ -110,8 +110,29 @@ def rewrite_eye(g, ops):
                 OpTypePattern("Const", name="fill_value"),
             ]), "*"
         ])
+    pattern7 = \
+        OpTypePattern("MatrixDiag", name="output_eye_matrix", inputs=[
+            OpTypePattern("Fill", inputs=[
+                OpTypePattern("Reshape", inputs=[
+                    OpTypePattern("Minimum|Cast", name="min_or_cast"),
+                    "*",
+                ]),
+                OpTypePattern("Const", name="fill_value"),
+            ])
+        ])
+    pattern8 = \
+        OpTypePattern("MatrixSetDiag", name="output_eye_matrix", inputs=[
+            OpTypePattern("Fill"),
+            OpTypePattern("Fill", inputs=[
+                OpTypePattern("Reshape", inputs=[
+                    OpTypePattern("Minimum|Cast", name="min_or_cast"),
+                    "*",
+                ]),
+                OpTypePattern("Const", name="fill_value"),
+            ])
+        ])
 
-    for pattern in [pattern1, pattern2, pattern3, pattern4, pattern5, pattern6]:
+    for pattern in [pattern1, pattern2, pattern3, pattern4, pattern5, pattern6, pattern7, pattern8]:
         matcher = GraphMatcher(pattern, allow_reorder=True)
         match_results = list(matcher.match_ops(ops))
         for match_result in match_results:
@@ -146,6 +167,6 @@ def rewrite_eye(g, ops):
             zero_matrix = g.make_node("ConstantOfShape", matrix_shape_int64.output)
 
             g.make_node("EyeLike", zero_matrix.output, attr={"dtype": output_dtypes[0]},
-                        name=old_output.name, shapes=output_shapes, dtypes=output_dtypes)
+                        name=old_output.name, shapes=output_shapes, dtypes=output_dtypes, outputs=old_output.output)
 
     return g.get_nodes()

tf2onnx/tflite_handlers/tfl_nn.py

@@ -70,7 +70,7 @@ class TflDepthwiseConv2D:
     @classmethod
     def to_tf(cls, ctx, node, **kwargs):
         separate_fused_activation_function(ctx, node)
-        # No need to change 'padding' attribute
+        # No need to change 'padding' or 'depth_multiplier' attributes
         stride_h = node.get_attr_int("stride_h")
         stride_w = node.get_attr_int("stride_w")
         dilation_w_factor = node.get_attr_int("dilation_w_factor")
@@ -81,7 +81,6 @@ def to_tf(cls, ctx, node, **kwargs):
         del node.attr["stride_w"]
         del node.attr["dilation_h_factor"]
         del node.attr["dilation_w_factor"]
-        del node.attr["depth_multiplier"]  # TODO: use this?
         transpose_node = ctx.insert_new_node_on_input(node, "Transpose", node.input[1], name=None, perm=[1, 2, 3, 0])
         transpose_node.skip_conversion = True
         node.set_attr("data_format", "NHWC")