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

@@ -52,7 +52,8 @@ def load(self):
 
         logging.info("Loading TensorFlow model '{}'".format(self.model))
         outputs = self.kwargs.get("outputs", None)
-        self._graph_def = self._graph_def_from_model(outputs)
+        tags = self.kwargs.get("tags", None)
+        self._graph_def = self._graph_def_from_model(outputs, 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."
@@ -78,7 +79,7 @@ def load(self):
         return program
 
     # @abstractmethod
-    def _graph_def_from_model(self, outputs=None):
+    def _graph_def_from_model(self, outputs=None, tags=None):
         """Load TensorFlow model into GraphDef. Overwrite for different TF versions."""
         pass
 
@@ -129,7 +130,7 @@ def __init__(self, model, debug=False, **kwargs):
         """
         TFLoader.__init__(self, model, debug, **kwargs)
 
-    def _graph_def_from_model(self, outputs=None):
+    def _graph_def_from_model(self, outputs=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"):
@@ -160,7 +161,7 @@ def _graph_def_from_model(self, outputs=None):
                 graph_def = self._from_tf_keras_model(self.model)
                 return self.extract_sub_graph(graph_def, outputs)
             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, outputs)
             else:
                 raise NotImplementedError(msg.format(self.model))

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

@@ -68,7 +68,7 @@ def __init__(self, model, debug=False, **kwargs):
         """
         TFLoader.__init__(self, model, debug, **kwargs)
 
-    def _graph_def_from_model(self, outputs=None):
+    def _graph_def_from_model(self, outputs=None, tags=None):
         """Overwrites TFLoader._graph_def_from_model()"""
         msg = (
             "Expected model format: [SavedModel | [concrete_function] | "
@@ -92,7 +92,7 @@ def _graph_def_from_model(self, outputs=None):
                 elif _os_path.isfile(self.model) and self.model.endswith(".h5"):
                     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: