extract common constants

Author: nbcsm

examples/custom_op_via_python.py

@@ -4,8 +4,7 @@
 import tensorflow as tf
 import tf2onnx
 from onnx import helper
-
-_TENSORFLOW_DOMAIN = "ai.onnx.converters.tensorflow"
+from tf2onnx import constants
 
 
 def print_handler(ctx, node, name, args):
@@ -14,7 +13,7 @@ def print_handler(ctx, node, name, args):
     # becomes:
     #   T output = Identity(T Input)
     node.type = "Identity"
-    node.domain = _TENSORFLOW_DOMAIN
+    node.domain = constants.DEFAULT_CUSTOM_OP_OPSET.domain
     del node.input[1:]
     return node
 
@@ -26,7 +25,7 @@ def print_handler(ctx, node, name, args):
     _ = tf.identity(x_, name="output")
     onnx_graph = tf2onnx.tfonnx.process_tf_graph(sess.graph,
                                                  custom_op_handlers={"Print": print_handler},
-                                                 extra_opset=[helper.make_opsetid(_TENSORFLOW_DOMAIN, 1)],
+                                                 extra_opset=[constants.DEFAULT_CUSTOM_OP_OPSET],
                                                  input_names=["input:0"],
                                                  output_names=["output:0"])
     model_proto = onnx_graph.make_model("test")

tests/common.py

@@ -10,8 +10,7 @@
 from collections import defaultdict
 
 from distutils.version import LooseVersion
-from tf2onnx import utils
-from tf2onnx.tfonnx import DEFAULT_TARGET, POSSIBLE_TARGETS
+from tf2onnx import constants, utils
 
 __all__ = ["TestConfig", "get_test_config", "unittest_main",
            "check_tf_min_version", "skip_tf_versions",
@@ -26,8 +25,8 @@ class TestConfig(object):
     def __init__(self):
         self.platform = sys.platform
         self.tf_version = self._get_tf_version()
-        self.opset = int(os.environ.get("TF2ONNX_TEST_OPSET", 7))
-        self.target = os.environ.get("TF2ONNX_TEST_TARGET", ",".join(DEFAULT_TARGET)).split(',')
+        self.opset = int(os.environ.get("TF2ONNX_TEST_OPSET", constants.PREFERRED_OPSET))
+        self.target = os.environ.get("TF2ONNX_TEST_TARGET", ",".join(constants.DEFAULT_TARGET)).split(',')
         self.backend = os.environ.get("TF2ONNX_TEST_BACKEND", "onnxruntime")
         self.backend_version = self._get_backend_version()
         self.is_debug_mode = False
@@ -83,7 +82,7 @@ def load():
                                 choices=["caffe2", "onnxmsrtnext", "onnxruntime"],
                                 help="backend to test against")
             parser.add_argument("--opset", type=int, default=config.opset, help="opset to test against")
-            parser.add_argument("--target", default=",".join(config.target), choices=POSSIBLE_TARGETS,
+            parser.add_argument("--target", default=",".join(config.target), choices=constants.POSSIBLE_TARGETS,
                                 help="target platform")
             parser.add_argument("--debug", help="output debugging information", action="store_true")
             parser.add_argument("--temp_dir", help="temp dir")

tests/test_graph.py

@@ -18,12 +18,11 @@
 from onnx import helper
 
 import tf2onnx
+from tf2onnx import constants
 from tf2onnx.graph_matcher import OpTypePattern, GraphMatcher
 from tf2onnx.tfonnx import process_tf_graph
 from common import get_test_config, unittest_main
 
-_TENSORFLOW_DOMAIN = "ai.onnx.converters.tensorflow"
-
 
 # pylint: disable=missing-docstring
 
@@ -336,7 +335,7 @@ def print_handler(ctx, node, name, args):  # pylint: disable=unused-argument
             # becomes:
             #   T output = Identity(T Input)
             self.assertEqual(node.type, "Identity")
-            node.domain = _TENSORFLOW_DOMAIN
+            node.domain = constants.DEFAULT_CUSTOM_OP_OPSET.domain
             self.assertEqual(args[0], "mode")
             del node.input[1:]
             return node
@@ -348,7 +347,7 @@ def print_handler(ctx, node, name, args):  # pylint: disable=unused-argument
             g = process_tf_graph(sess.graph,
                                  custom_op_handlers={"Print": (print_handler, ["Identity", "mode"])},
                                  opset=self.config.opset,
-                                 extra_opset=helper.make_opsetid(_TENSORFLOW_DOMAIN, 1))
+                                 extra_opset=[constants.DEFAULT_CUSTOM_OP_OPSET])
             self.assertEqual(
                 'digraph { input1 [op_type=Placeholder shape="[2, 3]"] Print [op_type=Identity] '
                 'output [op_type=Identity] input1:0 -> Print Print:0 -> output }',

tf2onnx/constants.py

@@ -0,0 +1,29 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT license.
+
+"""
+common constants
+"""
+
+from onnx import helper
+
+# Built-in supported domains
+ONNX_DOMAIN = ""
+AI_ONNX_ML_DOMAIN = "ai.onnx.ml"
+MICROSOFT_DOMAIN = "com.microsoft"
+
+# Default opset version for onnx domain
+PREFERRED_OPSET = 7
+
+# Default opset for custom ops
+DEFAULT_CUSTOM_OP_OPSET = helper.make_opsetid("ai.onnx.converters.tensorflow", 1)
+
+# Target for the generated onnx graph. It possible targets:
+# onnx-1.1 = onnx at v1.1 (winml in rs4 is based on this)
+# caffe2 = include some workarounds for caffe2 and winml
+TARGET_RS4 = "rs4"
+TARGET_RS5 = "rs5"
+TARGET_RS6 = "rs6"
+TARGET_CAFFE2 = "caffe2"
+POSSIBLE_TARGETS = [TARGET_RS4, TARGET_RS5, TARGET_RS6, TARGET_CAFFE2]
+DEFAULT_TARGET = []

tf2onnx/convert.py

@@ -13,12 +13,9 @@
 from onnx import helper
 import tensorflow as tf
 
-from tf2onnx import utils
-from tf2onnx import loader
+from tf2onnx import constants, loader, utils
 from tf2onnx.graph import GraphUtil
-from tf2onnx.tfonnx import process_tf_graph, tf_optimize, DEFAULT_TARGET, POSSIBLE_TARGETS
-
-_TENSORFLOW_DOMAIN = "ai.onnx.converters.tensorflow"
+from tf2onnx.tfonnx import process_tf_graph, tf_optimize
 
 
 # pylint: disable=unused-argument
@@ -36,7 +33,8 @@ def get_args():
     parser.add_argument("--outputs", help="model output_names")
     parser.add_argument("--opset", type=int, default=None, help="onnx opset to use")
     parser.add_argument("--custom-ops", help="list of custom ops")
-    parser.add_argument("--target", default=",".join(DEFAULT_TARGET), choices=POSSIBLE_TARGETS, help="target platform")
+    parser.add_argument("--target", default=",".join(constants.DEFAULT_TARGET), choices=constants.POSSIBLE_TARGETS,
+                        help="target platform")
     parser.add_argument("--continue_on_error", help="continue_on_error", action="store_true")
     parser.add_argument("--verbose", help="verbose output", action="store_true")
     parser.add_argument("--fold_const", help="enable tf constant_folding transformation before conversion",
@@ -69,7 +67,7 @@ def get_args():
 
 
 def default_custom_op_handler(ctx, node, name, args):
-    node.domain = _TENSORFLOW_DOMAIN
+    node.domain = constants.DEFAULT_CUSTOM_OP_OPSET.domain
     return node
 
 
@@ -83,7 +81,7 @@ def main():
     if args.custom_ops:
         # default custom ops for tensorflow-onnx are in the "tf" namespace
         custom_ops = {op: (default_custom_op_handler, []) for op in args.custom_ops.split(",")}
-        extra_opset = [helper.make_opsetid(_TENSORFLOW_DOMAIN, 1)]
+        extra_opset = [constants.DEFAULT_CUSTOM_OP_OPSET]
     else:
         custom_ops = {}
         extra_opset = None

tf2onnx/schemas.py

@@ -12,7 +12,7 @@
 from collections import defaultdict, OrderedDict
 from onnx import defs
 
-ONNX_DOMAIN = ""
+from . import constants
 
 
 class OnnxOpSchema(object):
@@ -97,9 +97,9 @@ def _parse_domain_opset_versions(schemas):
 _domain_opset_versions = _parse_domain_opset_versions(_schemas)
 
 
-def get_schema(name, max_inclusive_opset_version, domain=ONNX_DOMAIN):
+def get_schema(name, max_inclusive_opset_version, domain=None):
     """Get schema by name within specific version."""
-    domain = domain or ONNX_DOMAIN
+    domain = domain or constants.ONNX_DOMAIN
     domain_version_schema_map = _schemas[name]
     version_schema_map = domain_version_schema_map[domain]
     for version, schema in version_schema_map.items():
@@ -108,7 +108,7 @@ def get_schema(name, max_inclusive_opset_version, domain=ONNX_DOMAIN):
     return None
 
 
-def get_max_supported_opset_version(domain=ONNX_DOMAIN):
+def get_max_supported_opset_version(domain=None):
     """Get max supported opset version by current onnx package given a domain."""
-    domain = domain or ONNX_DOMAIN
+    domain = domain or constants.ONNX_DOMAIN
     return _domain_opset_versions.get(domain, None)

tf2onnx/tfonnx.py

@@ -21,7 +21,7 @@
 from tensorflow.tools.graph_transforms import TransformGraph
 
 import tf2onnx
-from tf2onnx import schemas, utils
+from tf2onnx import constants, schemas, utils
 from tf2onnx.function import *  # pylint: disable=wildcard-import
 from tf2onnx.graph import Graph
 from tf2onnx.graph_matcher import OpTypePattern, GraphMatcher
@@ -32,16 +32,6 @@
 logging.basicConfig(level=logging.INFO)
 log = logging.getLogger("tf2onnx")
 
-# Target for the generated onnx graph. It possible targets:
-# onnx-1.1 = onnx at v1.1 (winml in rs4 is based on this)
-# caffe2 = include some workarounds for caffe2 and winml
-TARGET_RS4 = "rs4"
-TARGET_RS5 = "rs5"
-TARGET_RS6 = "rs6"
-TARGET_CAFFE2 = "caffe2"
-POSSIBLE_TARGETS = [TARGET_RS4, TARGET_RS5, TARGET_RS6, TARGET_CAFFE2]
-DEFAULT_TARGET = []
-
 
 # pylint: disable=useless-return,broad-except,logging-not-lazy,unused-argument,missing-docstring
 # FIXME:
@@ -178,7 +168,7 @@ def broadcast_op(ctx, node, name, args):
         node.set_attr("broadcast", 1)
         # this works around shortcomings in the broadcasting code
         # of caffe2 and winml/rs4.
-        if ctx.is_target(TARGET_RS4):
+        if ctx.is_target(constants.TARGET_RS4):
             # in rs4 mul and add do not support scalar correctly
             if not shape0:
                 if node.inputs[0].is_const():
@@ -201,7 +191,7 @@ def broadcast_op7(ctx, node, name, args):
     if shape0 != shape1:
         # this works around shortcomings in the broadcasting code
         # of caffe2 and winml/rs4.
-        if ctx.is_target(TARGET_RS4):
+        if ctx.is_target(constants.TARGET_RS4):
             # in rs4 mul and add do not support scalar correctly
             if not shape0:
                 if node.inputs[0].is_const():
@@ -1027,7 +1017,7 @@ def stridedslice_op(ctx, node, name, args):
 
 
 def pow_op(ctx, node, name, args):
-    if ctx.is_target(TARGET_CAFFE2):
+    if ctx.is_target(constants.TARGET_CAFFE2):
         # workaround a bug in caffe2 pre Feb2018, pow(a, b) becomes np.exp(np.log(a) * b)
         node.type = "Log"
         b = node.input[1]
@@ -1300,7 +1290,8 @@ def onehot_op9(ctx, node, name, args):
     # in ONNX, op's schema is (input, depth, value, @int axis), meaning of "value" is [off-value, on-value]
     # onnxruntime only supports int64
     output_dtype = ctx.get_dtype(node.input[2])
-    if ctx.is_target(TARGET_RS6) and output_dtype not in [onnx_pb.TensorProto.INT64, onnx_pb.TensorProto.INT32]:
+    if ctx.is_target(constants.TARGET_RS6) \
+            and output_dtype not in [onnx_pb.TensorProto.INT64, onnx_pb.TensorProto.INT32]:
         log.warning("unsupported dtype in onnxruntime, onehot-9 can't be used directly")
         onehot_op(ctx, node, name, args)
         return
@@ -1315,21 +1306,21 @@ def onehot_op9(ctx, node, name, args):
     off_on_value = ctx.make_node("Concat", [off_value, on_value], attr={"axis": 0}).output[0]
 
     indices = node.input[0]
-    if ctx.is_target(TARGET_RS6) and ctx.get_dtype(indices) != onnx_pb.TensorProto.INT64:
+    if ctx.is_target(constants.TARGET_RS6) and ctx.get_dtype(indices) != onnx_pb.TensorProto.INT64:
         indices = ctx.make_node("Cast", [indices], attr={"to": onnx_pb.TensorProto.INT64}).output[0]
     node.input[0] = indices
 
-    if ctx.is_target(TARGET_RS6) and ctx.get_dtype(depth) != onnx_pb.TensorProto.INT64:
+    if ctx.is_target(constants.TARGET_RS6) and ctx.get_dtype(depth) != onnx_pb.TensorProto.INT64:
         depth = ctx.make_node("Cast", [depth], attr={"to": onnx_pb.TensorProto.INT64}).output[0]
     node.input[1] = depth
 
-    if ctx.is_target(TARGET_RS6) and output_dtype != onnx_pb.TensorProto.INT64:
+    if ctx.is_target(constants.TARGET_RS6) and output_dtype != onnx_pb.TensorProto.INT64:
         off_on_value = ctx.make_node("Cast", [off_on_value], attr={"to": onnx_pb.TensorProto.INT64}).output[0]
     node.input[2] = off_on_value
 
     del node.input[3]
 
-    if ctx.is_target(TARGET_RS6) and output_dtype != onnx_pb.TensorProto.INT64:
+    if ctx.is_target(constants.TARGET_RS6) and output_dtype != onnx_pb.TensorProto.INT64:
         new_node_name = utils.make_name("onehot_output")
         new_node = ctx.insert_new_node_on_output("Cast", node.output[0], new_node_name, to=output_dtype)
         ctx.set_dtype(new_node.output[0], output_dtype)
@@ -2505,7 +2496,7 @@ def process_tf_graph(tf_graph, continue_on_error=False, verbose=False, target=No
     if inputs_as_nchw is None:
         inputs_as_nchw = []
     if target is None:
-        target = DEFAULT_TARGET
+        target = constants.DEFAULT_TARGET
 
     onnx_nodes, op_cnt, attr_cnt, output_shapes, dtypes = tensorflow_to_onnx(tf_graph, shape_override)
 
@@ -2558,9 +2549,9 @@ def process_tf_graph(tf_graph, continue_on_error=False, verbose=False, target=No
 
     # post-processing rewriters
     late_rewriters = []
-    if TARGET_RS5 in target:
+    if constants.TARGET_RS5 in target:
         late_rewriters.append(rewrite_incomplete_type_support_rs5)
-    if TARGET_RS6 in target:
+    if constants.TARGET_RS6 in target:
         late_rewriters.append(rewrite_incomplete_type_support_rs6)
     if late_rewriters:
         run_rewriters(g, late_rewriters, continue_on_error)

tf2onnx/utils.py

@@ -20,6 +20,7 @@
 from google.protobuf import text_format
 import onnx
 from onnx import helper, onnx_pb, defs, numpy_helper
+from . import constants
 
 #
 #  mapping dtypes from tensorflow to onnx
@@ -174,12 +175,7 @@ def get_tf_tensor_data(tensor):
         data = tensor.bool_val
     elif tensor.string_val:
         data = tensor.string_val
-    elif tensor.dtype in [
-            tf.int32,
-            tf.int64,
-            tf.float32,
-            tf.float16
-    ]:
+    elif tensor.dtype in [tf.int32, tf.int64, tf.float32, tf.float16]:
         data = None
     else:
         raise ValueError('tensor data not supported')
@@ -249,16 +245,13 @@ def make_onnx_inputs_outputs(name, elem_type, shape, **kwargs):
     return helper.make_tensor_value_info(name, elem_type, make_onnx_shape(shape), **kwargs)
 
 
-PREFERRED_OPSET = 7
-
-
 def find_opset(opset):
     """Find opset."""
     if opset is None or opset == 0:
         opset = defs.onnx_opset_version()
-        if opset > PREFERRED_OPSET:
+        if opset > constants.PREFERRED_OPSET:
             # if we use a newer onnx opset than most runtimes support, default to the one most supported
-            opset = PREFERRED_OPSET
+            opset = constants.PREFERRED_OPSET
     return opset