Don't copy/split transpose unless further optimization is likely (#1679)

* Don't copy/split tranpose unless further optimization is likely

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

* bump ci

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

Author: TomWildenhain-Microsoft

tests/test_optimizers.py

@@ -1436,7 +1436,6 @@ def test_two_transposes_switch_with_mul(self, shape, perm_input, perm_output):
                                    model_proto, remaining_transpose_num=0)
 
     @parameterized.expand([
-        ((1, 6, 8), (8, 1, 6), [2, 0, 1], [1, 2, 0]),
         ((1, 6, 8, 9), (1, 8, 9, 6), [0, 2, 3, 1], [0, 3, 1, 2]),
         ((1, 6, 8, 9, 2), (1, 8, 9, 2, 6), [0, 2, 3, 4, 1], [0, 4, 1, 2, 3]),
     ])

tf2onnx/optimizer/transpose_optimizer.py

@@ -323,8 +323,9 @@ def _handle_nhwc_tranpose(self, trans):
                 op_handler = self._handler_map[p.type]
                 return op_handler(trans, p)
             return False
-        if out_nodes:
-            # move transpose into branches to let Transposes can be "handled" in each branch
+        if out_nodes and trans.get_attr_value("perm") in [NCHW_TO_NHWC, NCDHW_TO_NDHWC]:
+            # Move transpose into branches to let Transposes can be "handled" in each branch.
+            # This will add more transpose ops, so only do this if further optimization is likely (check perm).
             for n in out_nodes:
                 branch_trans = n.graph.make_node("Transpose", [trans.input[0]], attr=trans.get_onnx_attrs())
                 n.graph.replace_input(n, trans.output[0], branch_trans.output[0])