refine custom opsets loading

Author: nbcsm

tf2onnx/custom_opsets/__init__.py

@@ -3,3 +3,8 @@
 """ custom tf2onnx mapping functions. """
 
 from . import ms
+from tf2onnx import constants
+
+DOMAIN_OPSETS = {
+    constants.MICROSOFT_DOMAIN: ms.OPSETS
+}

tf2onnx/tfonnx.py

@@ -2334,9 +2334,9 @@ def tensorflow_onnx_mapping(g, continue_on_error, custom_op_handlers):
     # 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_opsets.ms.OPSETS:
+            opsets = custom_opsets.DOMAIN_OPSETS.get(extra_opset.domain, None)
+            if opsets is not None:
+                for target_opset, op_map in opsets:
                     if target_opset <= extra_opset.version:
                         ops_mapping.update(op_map)
             else: