support extra opset mapping

Author: nbcsm

tf2onnx/convert.py

@@ -31,8 +31,10 @@ def get_args():
     parser.add_argument("--output", help="output model file")
     parser.add_argument("--inputs", help="model input_names")
     parser.add_argument("--outputs", help="model output_names")
-    parser.add_argument("--opset", type=int, default=None, help="onnx opset to use")
+    parser.add_argument("--opset", type=int, default=None, help="opset version to use for onnx domain")
     parser.add_argument("--custom-ops", help="list of custom ops")
+    parser.add_argument("--extra_opset", default=None,
+                        help="extra opset with format like domain:version, e.g. com.microsoft:1")
     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")
@@ -63,6 +65,11 @@ def get_args():
     if args.target:
         args.target = args.target.split(",")
 
+    if args.extra_opset:
+        tokens = args.extra_opset.split(':')
+        if len(tokens) != 2:
+            raise ValueError("invalid extra_opset argument")
+        args.extra_opset = [helper.make_opsetid(tokens[0], int(tokens[1]))]
     return args
 
 
@@ -78,13 +85,12 @@ def main():
     # support unknown dimensions.
     utils.ONNX_UNKNOWN_DIMENSION = args.unknown_dim
 
+    extra_opset = args.extra_opset or []
+    custom_ops = {}
     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 = [constants.DEFAULT_CUSTOM_OP_OPSET]
-    else:
-        custom_ops = {}
-        extra_opset = None
+        extra_opset.append(constants.DEFAULT_CUSTOM_OP_OPSET)
 
     # get the frozen tensorflow model from graphdef, checkpoint or saved_model.
     if args.graphdef:

tf2onnx/tfonnx.py

@@ -21,7 +21,7 @@
 from tensorflow.tools.graph_transforms import TransformGraph
 
 import tf2onnx
-from tf2onnx import constants, schemas, utils
+from tf2onnx import constants, custom, schemas, utils
 from tf2onnx.function import *  # pylint: disable=wildcard-import
 from tf2onnx.graph import Graph
 from tf2onnx.graph_matcher import OpTypePattern, GraphMatcher
@@ -2322,13 +2322,28 @@ def tensorflow_onnx_mapping(g, continue_on_error, custom_op_handlers):
     mapped_op = collections.Counter()
     unmapped_op = collections.Counter()
 
-    # create ops mapping for the desired opset
+    # create ops mapping for the desired opsets
     ops_mapping = {}
+
+    # load mapping for onnx domain
     for target_opset, op_map in _OPSETS:
         if target_opset <= g.opset:
             ops_mapping.update(op_map)
 
-    # apply custom ops on top of the assembled opset. We can either completment the opset
+    # load mapping for known extra opsets
+    # order matters, later mapping overrides earlier's
+    if g.extra_opset is not None:
+        for extra_opset in g.extra_opset:
+            if extra_opset.domain == constants.MICROSOFT_DOMAIN:
+                # microsoft domain
+                for target_opset, op_map in custom.ms.OPSETS:
+                    if target_opset <= extra_opset.version:
+                        ops_mapping.update(op_map)
+            else:
+                # unknown opset, assume used in custom_op_handlers, skip it
+                pass
+
+    # apply custom ops on top of the assembled opset. We can either complement the opset
     # or override existing ops with a custom op.
     if custom_op_handlers is not None:
         custom_opset = {k: v for k, v in custom_op_handlers.items()}
@@ -2337,7 +2352,7 @@ def tensorflow_onnx_mapping(g, continue_on_error, custom_op_handlers):
     ops = [n for n in g.get_nodes()]
     for node in ops:
         if node.need_skip():
-            log.debug("explictly skip node " + node.name)
+            log.debug("explicitly skip node " + node.name)
             continue
 
         op = node.type