refine onnx_to_numpy conversion a bit

Author: pengwa

tf2onnx/function/sparse_softmax_cross_entropy_with_logits.py

@@ -29,7 +29,7 @@ def sparse_softmax_cross_entropy_with_logits_op(ctx, node, name, args):
     logit_dtype = ctx.get_dtype(logit_name)
     utils.make_sure(logit_dtype, "Dtype of {} is None".format(logit_name))
 
-    dtype = utils.ONNX_TO_NUMPY_DTYPE[logit_dtype]
+    dtype = utils.map_onnx_to_numpy_type(logit_dtype)
     eye = np.eye(depth).astype(dtype)
     const_name = utils.make_name("const_eye")
     const_eye = ctx.make_const(name=const_name, np_val=eye)
@@ -82,7 +82,7 @@ def sparse_softmax_cross_entropy_with_logits_op_by_gathernd(ctx, node, name, arg
     make_gathernd(ctx, log_softmax.output[0], indices_with_id.output[0], gathernd_output,
                   gathernd_name, logit_dtype)
     const_name = utils.make_name("const_negative_one")
-    const_negative_one = ctx.make_const(const_name, np.array(-1).astype(utils.ONNX_TO_NUMPY_DTYPE[logit_dtype]))
+    const_negative_one = ctx.make_const(const_name, np.array(-1).astype(utils.map_onnx_to_numpy_type(logit_dtype)))
     mul2 = ctx.make_node(op_type="Mul", inputs=[const_negative_one.output[0], gathernd_output])
     shapes = node.output_shapes
     dtypes = node.output_dtypes

tf2onnx/rewriter/rnn_utils.py

@@ -198,7 +198,7 @@ def get_weights_from_const_node(g, node):
 
     if temp and temp.type == 'Const':
         val = temp.get_tensor_value(as_list=False)
-        dtype = utils.ONNX_TO_NUMPY_DTYPE[g.get_dtype(temp.output[0])]
+        dtype = utils.map_onnx_to_numpy_type(g.get_dtype(temp.output[0]))
         log.debug("found weights %s", temp.name)
     else:
         log.debug("weight node seems not to be Const, skip, node name is %s", temp.name)

tf2onnx/rewriter/unit_rewriter_base.py

@@ -420,7 +420,7 @@ def _workaround_fill_ch_init_node(self, initializer_input_id, rnn_props):
             return None
 
         fill_val = node.inputs[1].get_tensor_value()
-        fill_val_dtype = utils.ONNX_TO_NUMPY_DTYPE[self.g.get_dtype(node.input[1])]
+        fill_val_dtype = utils.map_onnx_to_numpy_type(self.g.get_dtype(node.input[1]))
 
         # this must be int64, since Concat's input data type must be consistent.
         num_direction_node = self.g.make_const(utils.make_name("Const"), np.array([1], dtype=np.float32))

tf2onnx/tfonnx.py

@@ -604,7 +604,7 @@ def relu6_op(ctx, node, name, args):
     # since onnx does not have relu6, compose it with multiple ops.
     old_output = node.output[0]
     dtype = ctx.get_dtype(node.input[0])
-    dtype = utils.ONNX_TO_NUMPY_DTYPE[dtype] if dtype else np.float32
+    dtype = utils.map_onnx_to_numpy_type(dtype) if dtype else np.float32
     shape = ctx.get_shape(node.input[0])
     if -1 in shape:
         # if the shape has unknown dims we need to do something like this for opset < 8 (=no broadcast for min/max):
@@ -651,7 +651,7 @@ def relu6_op8(ctx, node, name, args):
     # since onnx does not have relu6, compose it with multiple ops.
     old_output = node.output[0]
     dtype = ctx.get_dtype(node.input[0])
-    dtype = utils.ONNX_TO_NUMPY_DTYPE[dtype] if dtype else np.float32
+    dtype = utils.map_onnx_to_numpy_type(dtype) if dtype else np.float32
     node.type = "Max"
     # const tensor 6
     six_name = utils.make_name(node.name)
@@ -687,7 +687,7 @@ def sign_op(ctx, node, name, args):
     utils.make_sure(node_dtype, "Dtype of {} is None".format(node.name))
     if node_dtype in [onnx_pb.TensorProto.COMPLEX64, onnx_pb.TensorProto.COMPLEX128]:
         raise ValueError("dtype " + node_dtype + " is not supported in onnx for now")
-    input_tensor_type = utils.ONNX_TO_NUMPY_DTYPE[node_dtype]
+    input_tensor_type = utils.map_onnx_to_numpy_type(node_dtype)
     zero_name = utils.make_name("{}_zero".format(node.name))
     ctx.make_const(zero_name, np.array(0, dtype=input_tensor_type))
     greater_node = ctx.make_node("Greater", [node.input[0], zero_name])
@@ -1280,7 +1280,7 @@ def fused_batchnorm_op7(ctx, node, name, args):
     scale_shape = ctx.get_shape(node.input[1])
     mean_shape = ctx.get_shape(node.input[3])
     var_shape = ctx.get_shape(node.input[4])
-    val_type = utils.ONNX_TO_NUMPY_DTYPE[ctx.get_dtype(node.input[1])]
+    val_type = utils.map_onnx_to_numpy_type(ctx.get_dtype(node.input[1]))
 
     if mean_shape != scale_shape:
         new_mean_value = np.array(np.resize(node.inputs[3].get_tensor_value(as_list=False), scale_shape),
@@ -1387,11 +1387,11 @@ def fill_op(ctx, node, name, args):
     # In onnx the value is an attribute so we need to fetch the value as const which
     # sooner or later will be a problem for tensorflow-onnx.
     # ConstantOfShape in onnxruntime only support int64, so insert cast op
-    input_dtype_is_int64 = utils.ONNX_TO_NUMPY_DTYPE[ctx.get_dtype(node.input[0])] == np.int64
+    input_dtype_is_int64 = utils.map_onnx_to_numpy_type(ctx.get_dtype(node.input[0])) == np.int64
     if not input_dtype_is_int64:
         cast_node = ctx.insert_new_node_on_input(node, "Cast", node.input[0], to=onnx_pb.TensorProto.INT64)
     dtype = ctx.get_dtype(node.output[0])
-    value = np.array([node.inputs[1].get_tensor_value()]).astype(utils.ONNX_TO_NUMPY_DTYPE[dtype])
+    value = np.array([node.inputs[1].get_tensor_value()]).astype(utils.map_onnx_to_numpy_type(dtype))
     value_proto = numpy_helper.from_array(value)
     node.set_attr("value", value_proto)
     del node.input[1]
@@ -1601,7 +1601,7 @@ def zeroslike_op(ctx, node, name, args):
     # when params "dtype" used, tf will call another op "Fill" instead, so Cast is not needed here.
     input_dtype = ctx.get_dtype(node.input[0])
     node_name = utils.make_name("zero")
-    const_zero = ctx.make_const(node_name, np.array(0).astype(utils.ONNX_TO_NUMPY_DTYPE[input_dtype]))
+    const_zero = ctx.make_const(node_name, np.array(0).astype(utils.map_onnx_to_numpy_type(input_dtype)))
     shapes = node.output_shapes
     dtypes = node.output_dtypes
     ctx.remove_node(name)
@@ -2051,15 +2051,15 @@ def rewrite_constant_fold(g, ops):
                     log.info("folding node type=%s, name=%s" % (op.type, op.name))
                     if op.type == "Cast":
                         dst = op.get_attr_int("to")
-                        np_type = tf2onnx.utils.ONNX_TO_NUMPY_DTYPE[dst]
+                        np_type = tf2onnx.utils.map_onnx_to_numpy_type(dst)
                         val = np.cast[np_type](*inputs)
                     elif op.type == "ConcatV2":
                         axis = inputs[-1]
                         values = inputs[:-1]
                         val = func(tuple(values), axis)
                     elif op.type == "ListDiff":
                         out_type = op.get_attr_int("out_idx")
-                        np_type = tf2onnx.utils.ONNX_TO_NUMPY_DTYPE[out_type]
+                        np_type = tf2onnx.utils.map_onnx_to_numpy_type(out_type)
                         val = func(*inputs)
                         val = val.astype(np_type)
                     elif op.type in ["Pack"]:
@@ -2068,7 +2068,7 @@ def rewrite_constant_fold(g, ops):
                         val = func(inputs, axis=axis)
                     elif op.type == "Range":
                         dtype = op.get_attr_int("Tidx")
-                        np_type = tf2onnx.utils.ONNX_TO_NUMPY_DTYPE[dtype]
+                        np_type = tf2onnx.utils.map_onnx_to_numpy_type(dtype)
                         val = func(*inputs, dtype=np_type)
                     else:
                         val = func(*inputs)

tf2onnx/utils.py

@@ -138,7 +138,7 @@ def tf_to_onnx_tensor(tensor, name=""):
         dims = [1]
     is_raw, data = get_tf_tensor_data(tensor)
     if not is_raw and len(data) == 1 and np.prod(dims) > 1:
-        batch_data = np.zeros(dims, dtype=ONNX_TO_NUMPY_DTYPE[new_type])
+        batch_data = np.zeros(dims, dtype=map_onnx_to_numpy_type(new_type))
         batch_data.fill(data[0])
         onnx_tensor = numpy_helper.from_array(batch_data, name=name)
     else:
@@ -207,6 +207,10 @@ def map_numpy_to_onnx_dtype(np_dtype):
     raise ValueError("unsupported dtype " + np_dtype + " for mapping")
 
 
+def map_onnx_to_numpy_type(onnx_type):
+    return ONNX_TO_NUMPY_DTYPE[onnx_type]
+
+
 def node_name(name):
     """Get node name without io#."""
     pos = name.find(":")
@@ -228,12 +232,6 @@ def port_name(name, nr=0):
     return name + ":" + str(nr)
 
 
-def make_onnx_identity(node_input, node_output, name=None):
-    if name is None:
-        name = make_name("identity")
-    return helper.make_node("Identity", [node_input], [node_output], name=name)
-
-
 def make_onnx_inputs_outputs(name, elem_type, shape, **kwargs):
     """Wrapper for creating onnx graph inputs or outputs
        name,  # type: Text