Merge pull request #748 from jignparm/jignparm/fix_pad_for_dynamic_input

Fix Pad op for dynamic input (opset 11).

Author: guschmue

tf2onnx/graph.py

@@ -1067,7 +1067,9 @@ def insert_new_node_on_input(self, node, op_type, input_name, name=None, domain=
         Args:
             node: we want to replace the input for this node
             op_type: type for new operation
-            input_name: the names of the outputs above us
+            input_name: the name(s) of the outputs above us
+                if scalar, new node placed above input_name
+                if list, new node placed above input_name[0]. list is inputs into new node
             name: the name of the new op
             kwargs: attributes of the new node
 
@@ -1077,9 +1079,12 @@ def insert_new_node_on_input(self, node, op_type, input_name, name=None, domain=
         if name is None:
             name = utils.make_name(node.name)
         new_output = port_name(name)
-        new_node = self.make_node(op_type, [input_name], attr=kwargs, outputs=[new_output], name=name, domain=domain)
+        if not isinstance(input_name, list):
+            input_name = [input_name]
+
+        new_node = self.make_node(op_type, input_name, attr=kwargs, outputs=[new_output], name=name, domain=domain)
         for i, n in enumerate(node.input):
-            if n == input_name:
+            if n == input_name[0]:
                 node.input[i] = new_output
                 break
         return new_node

tf2onnx/onnx_opset/generator.py

@@ -16,9 +16,9 @@
 from tf2onnx import utils
 from tf2onnx.handler import tf_op
 
-
 logger = logging.getLogger(__name__)
 
+
 # pylint: disable=unused-argument,missing-docstring
 
 @tf_op(["Const", "ConstV2"])
@@ -102,6 +102,15 @@ def version_9(cls, ctx, node, **kwargs):
         node.set_attr("value", value_proto)
         del node.input[1]
 
+    @classmethod
+    def version_11(cls, ctx, node, **kwargs):
+        # cls.version_7(ctx, node, **kwargs)
+        node.type = "Expand"
+        node.input = [node.input[1], node.input[0]]
+        # cast shape to int64 if needed
+        if ctx.get_dtype(node.input[1]) != onnx_pb.TensorProto.INT64:
+            ctx.insert_new_node_on_input(node, "Cast", node.input[1], to=onnx_pb.TensorProto.INT64)
+
 
 @tf_op("Multinomial")
 class Multinomial:

tf2onnx/onnx_opset/nn.py

@@ -314,7 +314,7 @@ def version_1(cls, ctx, node, **kwargs):
         k_h, k_w, k_input_channels, k_channel_multiplier = kernel_shape
         if k_input_channels < 1:
             raise ValueError("input channel must be positive")
-        k_output_channels = k_input_channels  * k_channel_multiplier
+        k_output_channels = k_input_channels * k_channel_multiplier
 
         node.set_attr("kernel_shape", [k_h, k_w])
         strides = conv_dims_attr(node, "strides")
@@ -476,13 +476,16 @@ def version_11(cls, ctx, node, **kwargs):
         if mode not in [None, "constant", "reflect"]:
             raise ValueError(mode + " pad mode is not supported")
 
-        pads = node.inputs[1].get_tensor_value()
-        pads = np.array(pads).transpose().flatten().astype(np.int64)
-        node.inputs[1].set_tensor_value(pads)
+        # pads must be int64.
+        if ctx.get_dtype(node.input[1]) != onnx_pb.TensorProto.INT64:
+            ctx.insert_new_node_on_input(node, "Cast", node.input[1], to=onnx_pb.TensorProto.INT64)
+        ctx.insert_new_node_on_input(node, "Transpose", node.input[1])
+        shape_const = ctx.make_const(utils.make_name(node.name), np.array([-1]).astype(np.int64))
+        ctx.insert_new_node_on_input(node, "Reshape", [node.input[1], shape_const.name])
 
         origin_dtype = ctx.get_dtype(node.output[0])
-        if origin_dtype not in [TensorProto.FLOAT16, TensorProto.FLOAT,
-                                TensorProto.DOUBLE]:
+        if origin_dtype not in [TensorProto.FLOAT, TensorProto.DOUBLE,
+                                TensorProto.INT32, TensorProto.INT64]:
             cast_node = ctx.insert_new_node_on_input(node, "Cast", node.input[0])
             cast_node.set_attr("to", TensorProto.FLOAT)
             ctx.set_dtype(cast_node.output[0], TensorProto.FLOAT)
@@ -739,8 +742,12 @@ def version_7(cls, ctx, node, **kwargs):
         # 2: "loop" to generate mask matrix: generate col or row of matrix one by one
         g = ctx.create_new_graph_with_same_config()
         node_name = utils.make_name("const_zero_bool")
-        const_zero_bool = ctx.make_const(name=node_name, np_val=np.array([[0]]).astype(np.bool))
-        ctx.set_dtype(const_zero_bool.output[0], onnx_pb.TensorProto.BOOL)
+        const_zero_bool = g.make_const(name=node_name, np_val=np.array([[0]]).astype(np.bool))
+        g.set_dtype(const_zero_bool.output[0], onnx_pb.TensorProto.BOOL)
+
+        g.add_graph_input("trip", onnx_pb.TensorProto.INT64, [])
+        g.add_graph_input("cond", onnx_pb.TensorProto.BOOL, [])
+        g.add_graph_input("line", onnx_pb.TensorProto.BOOL, [-1, -1])
 
         # shift right the line and add zero at the left.
         new_line = g.make_node(op_type="Concat", inputs=[const_zero_bool.output[0], "line"],
@@ -754,10 +761,6 @@ def version_7(cls, ctx, node, **kwargs):
         g.make_node("Identity", ["line"], outputs=["res"])
         g.make_node("Identity", [slice_node], outputs=["line_out"])
 
-        g.add_graph_input("trip", onnx_pb.TensorProto.INT64, [])
-        g.add_graph_input("cond", onnx_pb.TensorProto.BOOL, [])
-        g.add_graph_input("line", onnx_pb.TensorProto.BOOL, [-1, -1])
-
         g.add_graph_output("cond_out", onnx_pb.TensorProto.BOOL, [])
         g.add_graph_output("line_out", onnx_pb.TensorProto.BOOL, [-1, -1])
         g.add_graph_output("res", onnx_pb.TensorProto.BOOL, [-1, -1])