Merge pull request #92 from pengwa/fix-mirror-pad

fix the MirrorPad conversion failure

Author: guschmue

tests/test_graph.py

@@ -20,15 +20,20 @@
 _TENSORFLOW_DOMAIN = "ai.onnx.converters.tensorflow"
 
 
-def onnx_to_graphviz(g):
+def onnx_to_graphviz(g, include_attrs = False):
     g2 = gv.Digraph()
     for node in g.get_nodes():
         kwarg = {}
         attr = node.attr
-        if "shape" in attr:
-            kwarg["shape"] = str(attr["shape"].ints)
-        if "broadcast" in attr:
-            kwarg["broadcast"] = str(attr["broadcast"].i)
+        if include_attrs:
+            for a in attr:
+               kwarg[a] = str(helper.get_attribute_value(attr[a])) 
+        else: 
+            if "shape" in attr:
+                kwarg["shape"] = str(attr["shape"].ints)
+            if "broadcast" in attr:
+                kwarg["broadcast"] = str(attr["broadcast"].i)
+        
         g2.node(node.name, op_type=node.type, **kwarg)
     for node in g.get_nodes():
         for i in node.input:
@@ -287,6 +292,21 @@ def print_handler(ctx, node, name, args):
                 'digraph { Print [op_type=Identity] output [op_type=Identity] input1:0 -> Print Print:0 -> output }',
                 onnx_to_graphviz(g))
 
+    def test_pad(self):
+        with tf.Session() as sess:
+            t = tf.constant([[1, 2, 3], [4, 5, 6]], name= "input1")
+            paddings = tf.constant([[1, 1,], [2, 2]], name="paddings")
+            a = tf.pad(t, paddings, "CONSTANT", "const_no_val")
+            b = tf.pad(t, paddings, "CONSTANT", "const_with_val", 999)
+            c= tf.pad(t, paddings, "REFLECT", "reflect") 
+            g = process_tf_graph(sess.graph)
+
+            self.assertEqual('digraph { const_no_val [op_type=Pad pads="[1, 2, 1, 2]"]'
+                ' const_with_val [op_type=Pad pads="[1, 2, 1, 2]" value=999]'
+                ' reflect [mode="b\'reflect\'" op_type=Pad pads="[1, 2, 1, 2]"]'
+                ' input1:0 -> const_no_val input1:0 -> const_with_val input1:0 -> reflect }', 
+                onnx_to_graphviz(g, True))
+
 
 if __name__ == '__main__':
     unittest.main()

tf2onnx/tfonnx.py

@@ -710,9 +710,23 @@ def splitv_op(ctx, node, name, args):
 
 
 def pad_op(ctx, node, name, args):
-    # T output = Pad(T input, Tpaddings paddings, @type Tpaddings)
+    # T output = Pad(T input, int32 paddings, @type Tpaddings), CONST model using default value
+    #  or PadV2(T input, int32 paddings, T constant_value, @type Tpaddings), CONST mode - default value specified
+    #  or MirrorPad(T input, int32 paddings, @type Tpaddings, @STRING mode), other mode.
     # T output = Pad(T data, @STRING mode, @INTS pads, @FLOAT value)
     paddings = np.array(node.inputs[1].get_tensor_value()).transpose().flatten()
+    mode = node.get_attr("mode")
+    if mode:
+        mode = mode.s.decode("utf-8").lower()
+        node.set_attr("mode", mode)
+    if mode not in [None, "constant", "reflect"]:
+        raise ValueError(mode + " pad mode is not supported")
+    
+    if mode in [None, "constant"] and len(node.input) == 3:
+        const_val = node.inputs[2].get_tensor_value()[0]
+        node.set_attr("value", const_val)
+        ctx.remove_input(node, node.input[2])
+       
     ctx.remove_input(node, node.input[1])
     node.set_attr("pads", paddings)
     return node
@@ -1063,11 +1077,13 @@ def fused_batchnorm_op7(ctx, node, name, args):
     "Mean": (reduce_op, ["ReduceMean"]),
     "Min": (reduce_op, ["ReduceMin"]),
     "Minimum": (minmax_op, ["Min"]),
+    "MirrorPad": (pad_op, ["Pad"]),
     "Mul": (broadcast_op, []),
     "Neg": (direct_op, []),
     "NoOp": (no_op, []),
     "NotEqual": (direct_op, ["Not"]),
     "Pad": (pad_op, []),
+    "PadV2": (pad_op, ["Pad"]),
     "Placeholder": (placeholder_op, []),
     "PlaceholderV2": (placeholder_op, []),
     "PlaceholderWithDefault": (placeholder_op, []),