passing tags for use when loading TF saved_model

Some saved_models have more than one tag, e.g.,
MobileBERT SQuAD 1.1 checkpoints. Need a way
to specify tag

Author: freedomtan

coremltools/converters/mil/frontend/tensorflow/load.py

@@ -62,7 +62,8 @@ def load(self):
         logging.info("Loading TensorFlow model '{}'".format(self.model))
         outputs = self.kwargs.get("outputs", None)
         output_names = get_output_names(outputs)
-        self._graph_def = self._graph_def_from_model(output_names)
+        tags = self.kwargs.get("tags", None)
+        self._graph_def = self._graph_def_from_model(output_names, tags)
 
         if self._graph_def is not None and len(self._graph_def.node) == 0:
             msg = "tf.Graph should have at least 1 node, Got empty graph."
@@ -88,7 +89,7 @@ def load(self):
         return program
 
     # @abstractmethod
-    def _graph_def_from_model(self, output_names=None):
+    def _graph_def_from_model(self, output_names=None, tags=None):
         """Load TensorFlow model into GraphDef. Overwrite for different TF versions."""
         pass
 
@@ -139,7 +140,7 @@ def __init__(self, model, debug=False, **kwargs):
         """
         TFLoader.__init__(self, model, debug, **kwargs)
 
-    def _graph_def_from_model(self, output_names=None):
+    def _graph_def_from_model(self, output_names=None, tags=None):
         """Overwrites TFLoader._graph_def_from_model()"""
         msg = "Expected model format: [tf.Graph | .pb | SavedModel | tf.keras.Model | .h5], got {}"
         if isinstance(self.model, tf.Graph) and hasattr(self.model, "as_graph_def"):
@@ -170,7 +171,7 @@ def _graph_def_from_model(self, output_names=None):
                 graph_def = self._from_tf_keras_model(self.model)
                 return self.extract_sub_graph(graph_def, output_names)
             elif os.path.isdir(str(self.model)):
-                graph_def = self._from_saved_model(self.model)
+                graph_def = self._from_saved_model(self.model, tags=tags)
                 return self.extract_sub_graph(graph_def, output_names)
             else:
                 raise NotImplementedError(msg.format(self.model))

coremltools/converters/mil/frontend/tensorflow2/load.py

@@ -96,7 +96,7 @@ def __init__(self, model, debug=False, **kwargs):
             fuse_dilation_conv,
         ]
 
-    def _get_concrete_functions_and_graph_def(self):
+    def _get_concrete_functions_and_graph_def(self, tags=None):
         msg = (
             "Expected model format: [SavedModel | [concrete_function] | "
             "tf.keras.Model | .h5], got {}"
@@ -120,7 +120,7 @@ def _get_concrete_functions_and_graph_def(self):
                      and (self.model.endswith(".h5") or self.model.endswith(".hdf5")):
                     cfs = self._concrete_fn_from_tf_keras_or_h5(self.model)
                 elif _os_path.isdir(self.model):
-                    saved_model = _tf.saved_model.load(self.model)
+                    saved_model = _tf.saved_model.load(self.model, tags=tags)
                     sv = saved_model.signatures.values()
                     cfs = sv if isinstance(sv, list) else list(sv)
                 else:
@@ -132,9 +132,9 @@ def _get_concrete_functions_and_graph_def(self):
 
         return cfs, graph_def
 
-    def _graph_def_from_model(self, output_names=None):
+    def _graph_def_from_model(self, output_names=None, tags=None):
         """Overwrites TFLoader._graph_def_from_model()"""
-        _, graph_def = self._get_concrete_functions_and_graph_def()
+        _, graph_def = self._get_concrete_functions_and_graph_def(tags=tags)
         return self.extract_sub_graph(graph_def, output_names)
 
     def _tf_ssa_from_graph_def(self, fn_name="main"):