Optimize Transpose->Pad when 'pads' input is not a constant (#1331)

When the second input is not a constant, let's shuffle it with Split followed by Concat.
There are examples of models, where this non-constant input gets constant folded anyway by a framework.
Even if that's not the case. Split+Concat of 8 (or 10) element tensor should be a good trade for Transpose pair.

Signed-off-by: Mateusz Tabaka <[email protected]>

Author: mateusztabaka

tests/test_optimizers.py

@@ -922,6 +922,34 @@ def test_transpose_pad11(self, input_shape, output_shape, pads, perm_input, perm
         self.run_transpose_compare(["res"], {"X": np.random.randn(*input_shape).astype(np.float32)},
                                    model_proto, remaining_transpose_num=0)
 
+    @parameterized.expand([
+        ((1, 3, 4, 5), (2, 6, 4, 8), [1, 0, 1, 3, 0, 0, 2, 0], [0, 2, 3, 1], [0, 3, 1, 2]),
+        ((1, 3, 4, 5, 6), (2, 5, 6, 8, 10), [1, 0, 1, 3, 1, 0, 2, 2, 1, 1], [0, 2, 3, 4, 1], [0, 4, 1, 2, 3]),
+    ])
+    @check_opset_min_version(11, "pad")
+    def test_transpose_pad11_non_const_pads(self, input_shape, output_shape, pads, perm_input, perm_output):
+
+        pads_val = np.array(pads, dtype=np.int64)
+        node0 = helper.make_node("Transpose", ["X"], ["Y"], perm=perm_input, name="trans_1")
+        node1 = helper.make_node("Pad", ["Y", "Pads"], ["Z"], name="pad")
+        node2 = helper.make_node("Transpose", ["Z"], ["res"], perm=perm_output, name="trans_2")
+
+        graph = helper.make_graph(
+            [node0, node1, node2],
+            "transpose-pad-test",
+            [helper.make_tensor_value_info("X", TensorProto.FLOAT, input_shape),
+             helper.make_tensor_value_info("Pads", TensorProto.INT64, pads_val.shape)],
+            [helper.make_tensor_value_info("res", TensorProto.FLOAT, output_shape)],
+        )
+
+        model_proto = self.make_model(graph, producer_name="onnx-tests")
+        self.run_transpose_compare(["res"],
+                                   {
+                                       "X": np.random.randn(*input_shape).astype(np.float32),
+                                       "Pads": pads_val
+                                   },
+                                   model_proto, remaining_transpose_num=0)
+
     @parameterized.expand([
         ((1, 3, 4, 5), (1, 3, 1, 1), [0, 2, 3, 1], [0, 3, 1, 2]),
         ((1, 3, 4, 5, 6), (1, 3, 1, 1, 1), [0, 2, 3, 4, 1], [0, 4, 1, 2, 3]),

tf2onnx/optimizer/transpose_optimizer.py

@@ -676,7 +676,21 @@ def _pad_handler(self, trans, node):
                 input1.set_tensor_value(new_pads)
                 input1.data_format = "NCHW"
             return self._switch_transpose_and_node(node, trans)
-        return False
+        # when the second input is not a constant, let's shuffle it with Split followed by Concat
+        # there are examples of models, where this non-constant input
+        # gets constant folded anyway by a framework.
+        split = self._g.make_node("Split", inputs=[node.input[1]], attr={}, output_count=trans_rank * 2)
+        pads = split.output
+        if trans_rank == 4:
+            new_pads = self._g.make_node("Concat", [pads[0], pads[3], pads[1], pads[2],
+                                                    pads[4], pads[7], pads[5], pads[6]],
+                                         {'axis': 0})
+        else:
+            new_pads = self._g.make_node("Concat", [pads[0], pads[4], pads[1], pads[2], pads[3],
+                                                    pads[5], pads[9], pads[6], pads[7], pads[8]],
+                                         {'axis': 0})
+        self._g.replace_input(node, node.input[1], new_pads.output[0], 1)
+        return self._switch_transpose_and_node(node, trans)
 
     def _reducemean_handler(self, trans, node):
         axes = node.get_attr("axes").ints