Merge branch 'master' of https://github.com/onnx/tensorflow-onnx into input

Author: sdpython

tests/test_optimizers.py

@@ -912,7 +912,6 @@ def test_duplicated_duplicated_attributes(self):
                                                 op_type="ReduceSum", remaining_op_num=2)
 
     def _check_initializer_num(self, graph_proto, num):
-        print(len(graph_proto.initializer))
         return num == len(graph_proto.initializer)
 
     def test_duplicated_duplicated_constant(self):

tf2onnx/graph_matcher.py

@@ -50,6 +50,7 @@ def __init__(self, op_type, name=None, inputs=None):
             input_pattern if isinstance(input_pattern, OpTypePattern) else
             OpTypePattern(input_pattern) for input_pattern in inputs
         ]
+        self.op_type_set = set(op_type.split('|')) if op_type else set()
 
     @property
     def op_type(self):
@@ -154,7 +155,7 @@ def _is_op_type_same(op, pattern):
         if pattern.op_type == "*":
             return True
 
-        if op.type in pattern.op_type.split('|'):
+        if op.type in pattern.op_type_set:
             return True
 
         return False

tf2onnx/onnx_opset/controlflow.py

@@ -11,7 +11,6 @@
 
 import copy
 import logging
-import sys
 
 import numpy as np
 
@@ -524,7 +523,7 @@ def version_7(cls, ctx, node, **kwargs):
         maximum_iterations_name = node.input[1]
         maximum_iterations = node.inputs[1].get_tensor_value()
         if maximum_iterations == -1:
-            maximum_iterations = sys.maxsize
+            maximum_iterations = np.iinfo(np.int64).max
         consumers = ctx.find_output_consumers(maximum_iterations_name)
         external_consumers = [c for c in consumers if c != node and c.type != 'TensorListReserve']
         if len(external_consumers) == 0:

tf2onnx/onnx_opset/math.py

@@ -183,7 +183,7 @@ def version_8(cls, ctx, node, **kwargs):
 
     @classmethod
     def version_12(cls, ctx, node, **kwargs):
-        node.name = 'Clip' # clip supports all types now
+        node.type = 'Clip' # clip supports all types now
 
 @tf_op("Softmax")
 class Softmax:
@@ -545,9 +545,9 @@ def version_11(cls, ctx, node, **kwargs):
                              shapes=shapes, dtypes=dtypes, domain=constants.ONNX_DOMAIN, attr={'direction': direction})
 
         if node.maybe_cast_input([supported, supported], type_map):
-            cast_back_node = ctx.insert_new_node_on_output("Cast", node.output[0],
-                                                           name=utils.make_name(node.name) + "_castback")
-            cast_back_node.set_attr("to", dtypes[0])
+            cast_back_node = ctx.insert_new_node_on_output(
+                "Cast", node.output[0], name=utils.make_name(node.name) + "_castback",
+                to=dtypes[0])
             ctx.set_dtype(cast_back_node.output[0], dtypes[0])
             ctx.copy_shape(node.name, cast_back_node.output[0])
 

tf2onnx/onnx_opset/nn.py

@@ -637,14 +637,13 @@ def version_1(cls, ctx, node, **kwargs):
         origin_dtype = ctx.get_dtype(node.output[0])
         if origin_dtype not in [onnx_pb.TensorProto.FLOAT16, onnx_pb.TensorProto.FLOAT,
                                 onnx_pb.TensorProto.DOUBLE]:
-            cast_node = ctx.insert_new_node_on_input(node, "Cast", node.input[0])
-            cast_node.set_attr("to", onnx_pb.TensorProto.FLOAT)
+            cast_node = ctx.insert_new_node_on_input(node, "Cast", node.input[0], to=onnx_pb.TensorProto.FLOAT)
             ctx.set_dtype(cast_node.output[0], onnx_pb.TensorProto.FLOAT)
             ctx.copy_shape(node.name, cast_node.output[0])
 
             cast_back_node = ctx.insert_new_node_on_output("Cast", node.output[0],
-                                                           name=utils.make_name(node.name) + "_castback")
-            cast_back_node.set_attr("to", origin_dtype)
+                                                           name=utils.make_name(node.name) + "_castback",
+                                                           to=origin_dtype)
             ctx.set_dtype(cast_back_node.output[0], origin_dtype)
             ctx.copy_shape(node.name, cast_back_node.output[0])
 
@@ -667,14 +666,13 @@ def version_11(cls, ctx, node, **kwargs):
         origin_dtype = ctx.get_dtype(node.output[0])
         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)
+            cast_node = ctx.insert_new_node_on_input(node, "Cast", node.input[0], to=TensorProto.FLOAT)
             ctx.set_dtype(cast_node.output[0], TensorProto.FLOAT)
             ctx.copy_shape(node.name, cast_node.output[0])
 
             cast_back_node = ctx.insert_new_node_on_output("Cast", node.output[0],
-                                                           name=utils.make_name(node.name) + "_castback")
-            cast_back_node.set_attr("to", origin_dtype)
+                                                           name=utils.make_name(node.name) + "_castback",
+                                                           to=origin_dtype)
             ctx.set_dtype(cast_back_node.output[0], origin_dtype)
             ctx.copy_shape(node.name, cast_back_node.output[0])
 

tf2onnx/onnx_opset/tensor.py

@@ -31,8 +31,7 @@ def _convert_shapenode_to_int64(ctx, node, input_number):
     """cast int32 shape into int64 shape."""
     name = node.input[input_number]
 
-    cast_node = ctx.insert_new_node_on_input(node, "Cast", name)
-    cast_node.set_attr("to", onnx_pb.TensorProto.INT64)
+    cast_node = ctx.insert_new_node_on_input(node, "Cast", name, to=onnx_pb.TensorProto.INT64)
     ctx.set_dtype(cast_node.output[0], onnx_pb.TensorProto.INT64)
     ctx.copy_shape(name, cast_node.output[0])
 
@@ -46,14 +45,14 @@ def _wrap_concat_with_cast(ctx, node):
         output_name = node.output[0]
         # cast each inputs to float
         for i, inp in enumerate(node.inputs):
-            input_cast = ctx.insert_new_node_on_input(node, "Cast", node.input[i])
-            input_cast.set_attr("to", onnx_pb.TensorProto.FLOAT)
+            input_cast = ctx.insert_new_node_on_input(node, "Cast", node.input[i],
+                                                      to=onnx_pb.TensorProto.FLOAT)
             ctx.set_dtype(input_cast.output[0], onnx_pb.TensorProto.FLOAT)
         next_nodes = ctx.find_output_consumers(node.output[0])
         # cast output back to dtype unless the next op is a cast
         if next_nodes[0].type != "Cast":
-            output_cast = ctx.insert_new_node_on_output("Cast", output_name, name=node.child_name())
-            output_cast.set_attr("to", dtype)
+            output_cast = ctx.insert_new_node_on_output("Cast", output_name, name=node.child_name(),
+                                                        to=dtype)
             ctx.set_dtype(output_cast.output[0], dtype)
             ctx.copy_shape(output_name, output_cast.output[0])
 
@@ -157,15 +156,14 @@ def version_5(cls, ctx, node, **kwargs):
             return
 
         # onnx < opset 8 does not know reshape for other types than float*, wrap the reshape in casts
-        input_cast = ctx.insert_new_node_on_input(node, "Cast", node.input[0])
-        input_cast.set_attr("to", onnx_pb.TensorProto.FLOAT)
+        input_cast = ctx.insert_new_node_on_input(node, "Cast", node.input[0], to=onnx_pb.TensorProto.FLOAT)
         ctx.copy_shape(node.output[0], input_cast.output[0])
 
         # if the next node is already a cast we don't need to insert another one
         next_nodes = ctx.find_output_consumers(node.output[0])
         if len(next_nodes) != 1 or next_nodes[0].type != "Cast":
-            output_cast = ctx.insert_new_node_on_output("Cast", node.output[0], name=node.child_name())
-            output_cast.set_attr("to", dtype)
+            output_cast = ctx.insert_new_node_on_output("Cast", node.output[0], name=node.child_name(),
+                                                        to=dtype)
             ctx.set_dtype(output_cast.output[0], dtype)
             ctx.copy_shape(node.output[0], output_cast.output[0])
 
@@ -742,16 +740,17 @@ def version_1(cls, ctx, node, **kwargs):
                 if node.inputs[0].type == "Cast" and len(ctx.find_output_consumers(node.inputs[0].output[0])) == 1:
                     # override the previous cast
                     cast_node = node.inputs[0]
+                    cast_node.set_attr("to", onnx_pb.TensorProto.FLOAT)
                 else:
-                    cast_node = ctx.insert_new_node_on_input(node, "Cast", node.input[0])
+                    cast_node = ctx.insert_new_node_on_input(node, "Cast", node.input[0],
+                                                             to=onnx_pb.TensorProto.FLOAT)
                     nodes.insert(0, cast_node)
-                cast_node.set_attr("to", onnx_pb.TensorProto.FLOAT)
                 ctx.set_dtype(cast_node.output[0], onnx_pb.TensorProto.FLOAT)
                 ctx.copy_shape(node.input[0], cast_node.output[0])
                 # undo the cast afer slice
                 name = utils.make_name(node.name)
-                cast_node = ctx.insert_new_node_on_output("Cast", nodes[-1].output[0], name)
-                cast_node.set_attr("to", input_dtype)
+                cast_node = ctx.insert_new_node_on_output("Cast", nodes[-1].output[0], name,
+                                                          to=input_dtype)
                 ctx.set_dtype(cast_node.output[0], input_dtype)
                 ctx.copy_shape(node.output[0], cast_node.output[0])
                 nodes.append(cast_node)
@@ -1180,8 +1179,7 @@ def version_1(cls, ctx, node, **kwargs):
         if dtype == onnx_pb.TensorProto.INT64:
             return
         op_name = utils.make_name(node.name)
-        output_cast = ctx.insert_new_node_on_output("Cast", node.output[0], name=op_name)
-        output_cast.set_attr("to", dtype)
+        output_cast = ctx.insert_new_node_on_output("Cast", node.output[0], name=op_name, to=dtype)
         ctx.set_dtype(output_cast.output[0], dtype)
         ctx.copy_shape(node.output[0], output_cast.output[0])
 
@@ -1555,8 +1553,7 @@ def version_8(cls, ctx, node, **kwargs):
 
         seq_len_dtype = ctx.get_dtype(node.input[1])
         if seq_len_dtype != onnx_pb.TensorProto.INT64:
-            cast_node = ctx.insert_new_node_on_input(node, "Cast", node.input[1])
-            cast_node.set_attr("to", onnx_pb.TensorProto.INT64)
+            cast_node = ctx.insert_new_node_on_input(node, "Cast", node.input[1], to=onnx_pb.TensorProto.INT64)
             ctx.set_dtype(cast_node.output[0], onnx_pb.TensorProto.INT64)
             ctx.copy_shape(node.input[1], cast_node.output[0])
 
@@ -1762,8 +1759,8 @@ def version_11(cls, ctx, node, **kwargs):
             # cast to int64 if needed
             if dtypes[1] != onnx_pb.TensorProto.UINT64:
                 cast_node = ctx.insert_new_node_on_output("Cast", node.output[1],
-                                                          name=utils.make_name(node.name) + "_cast")
-                cast_node.set_attr("to", dtypes[1])
+                                                          name=utils.make_name(node.name) + "_cast",
+                                                          to=dtypes[1])
                 ctx.set_dtype(cast_node.output[0], dtypes[1])
                 ctx.copy_shape(node.output[1], cast_node.output[0])
             # FIXME: the indices in onnx are not the same as in tensorflow.

tf2onnx/tf_utils.py

@@ -13,7 +13,6 @@
 from distutils.version import LooseVersion
 
 import numpy as np
-import six
 import tensorflow as tf
 
 from tensorflow.core.framework import types_pb2, tensor_pb2
@@ -70,7 +69,7 @@ def get_tf_tensor_data(tensor):
     """Get data from tensor."""
     make_sure(isinstance(tensor, tensor_pb2.TensorProto), "Require TensorProto")
     np_data = tensor_util.MakeNdarray(tensor)
-    make_sure(isinstance(np_data, np.ndarray), "{} isn't ndarray".format(np_data))
+    make_sure(isinstance(np_data, np.ndarray), "%r isn't ndarray", np_data)
     return np_data
 
 
@@ -83,7 +82,7 @@ def get_tf_const_value(op, as_list=True):
         when as_list=False, return np.array(1), type is <class 'numpy.ndarray'>
         when as_list=True, return 1, type is <class 'int'>.
     """
-    make_sure(is_tf_const_op(op), "{} isn't a const op".format(op.name))
+    make_sure(is_tf_const_op(op), "%r isn't a const op", op.name)
     value = get_tf_tensor_data(op.get_attr("value"))
     if as_list:
         value = value.tolist()
@@ -119,9 +118,6 @@ def map_tf_dtype(dtype):
 
 def get_tf_node_attr(node, name):
     """Parser TF node attribute."""
-    if six.PY2:
-        # For python2, TF get_attr does not accept unicode
-        name = str(name)
     return node.get_attr(name)
 
 
@@ -136,14 +132,14 @@ def tflist_to_onnx(g, shape_override):
     """
 
     # ignore the following attributes
-    ignored_attr = ["unknown_rank", "_class", "Tshape", "use_cudnn_on_gpu", "Index", "Tpaddings",
+    ignored_attr = {"unknown_rank", "_class", "Tshape", "use_cudnn_on_gpu", "Index", "Tpaddings",
                     "TI", "Tparams", "Tindices", "Tlen", "Tdim", "Tin", "dynamic_size", "Tmultiples",
                     "Tblock_shape", "Tcrops", "index_type", "Taxis", "U", "maxval",
                     "Tout", "Tlabels", "Tindex", "element_shape", "Targmax", "Tperm", "Tcond",
                     "T_threshold", "element_dtype", "shape_type", "_lower_using_switch_merge",
                     "parallel_iterations", "_num_original_outputs", "output_types", "output_shapes",
                     "key_dtype", "value_dtype", "Tin", "Tout", "capacity", "component_types", "shapes",
-                    "Toutput_types"]
+                    "Toutput_types"}
 
     node_list = g.get_operations()
     functions = {}
@@ -176,12 +172,11 @@ def tflist_to_onnx(g, shape_override):
             attr_cnt[a] += 1
             if a == "dtype":
                 attr[a] = map_tf_dtype(get_tf_node_attr(node, "dtype"))
-            elif a in ["T"]:
+            elif a == "T":
                 dtype = get_tf_node_attr(node, a)
-                if dtype:
-                    if not isinstance(dtype, list):
-                        dtypes[node.name] = map_tf_dtype(dtype)
-            elif a in ["output_type", "output_dtype", "out_type", "Tidx", "out_idx"]:
+                if dtype and not isinstance(dtype, list):
+                    dtypes[node.name] = map_tf_dtype(dtype)
+            elif a in {"output_type", "output_dtype", "out_type", "Tidx", "out_idx"}:
                 # Tidx is used by Range
                 # out_idx is used by ListDiff
                 attr[a] = map_tf_dtype(get_tf_node_attr(node, a))
@@ -192,7 +187,7 @@ def tflist_to_onnx(g, shape_override):
             elif a == "output_shapes":
                 # we should not need it since we pull the shapes above already
                 pass
-            elif a in ["body", "cond", "then_branch", "else_branch"]:
+            elif a in {"body", "cond", "then_branch", "else_branch"}:
                 input_shapes = [inp.get_shape() for inp in node.inputs]
                 nattr = get_tf_node_attr(node, a)
                 attr[a] = nattr.name

tf2onnx/tfonnx.py

@@ -160,8 +160,8 @@ def rewrite_incomplete_type_support(g, ops, impacted_ops):
                         input_node.set_attr("to", onnx_pb.TensorProto.FLOAT)
                         g.set_dtype(input_name, onnx_pb.TensorProto.FLOAT)
                     else:
-                        cast_node = g.insert_new_node_on_input(op, "Cast", input_name)
-                        cast_node.set_attr("to", onnx_pb.TensorProto.FLOAT)
+                        cast_node = g.insert_new_node_on_input(op, "Cast", input_name,
+                                                               to=onnx_pb.TensorProto.FLOAT)
                         g.set_dtype(cast_node.output[0], onnx_pb.TensorProto.FLOAT)
                         g.copy_shape(input_name, cast_node.output[0])
                         cast_inserted.append(cast_node)
@@ -171,8 +171,8 @@ def rewrite_incomplete_type_support(g, ops, impacted_ops):
                     name = utils.make_name(op.name)
                     logger.debug("insert cast back for node %s on output %s [dtype=%s]", op.name, output_name,
                                  output_dtype)
-                    output_cast = g.insert_new_node_on_output("Cast", output_name, name=name)
-                    output_cast.set_attr("to", output_dtype)
+                    output_cast = g.insert_new_node_on_output("Cast", output_name, name=name,
+                                                              to=output_dtype)
                     g.set_dtype(output_cast.output[0], output_dtype)
                     g.copy_shape(output_name, output_cast.output[0])
                     cast_inserted.append(output_cast)