passing tags for use when loading TF saved_model

[freedomtan]

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

[aseemw]

• I don't think its good idea to pass additional flags via kwargs. If its something to be officially it should be part of the argument list
• can you please add unit tests, its not clear to me what creates these tags in the first place.

[heydavid525]

We ned to add this to our public doc as it's part of converter API (whether in kwargs or not), with example use cases. Like @aseemw suggested, can you add a test which would serve as an example?

[freedomtan]

Let's back to what the problem is and why tags is needed first.

When I was trying to convert MobileBERT with

import coremltools as ct

mobilebert_model = ct.convert('mobilebert_squad_savedmodels/float/', source="tensorflow")

I got

Traceback (most recent call last):
  File "m.py", line 3, in <module>
    mobilebert_model = ct.convert('mobilebert_squad_savedmodels/float/', source="tensorflow", tags='serve')
  File "/home/freedom/tf-master/lib/python3.8/site-packages/coremltools/converters/_converters_entry.py", line 176, in convert
    mlmodel = mil_convert(
  File "/home/freedom/tf-master/lib/python3.8/site-packages/coremltools/converters/mil/converter.py", line 128, in mil_convert
    proto = mil_convert_to_proto(model, convert_from, convert_to,
  File "/home/freedom/tf-master/lib/python3.8/site-packages/coremltools/converters/mil/converter.py", line 171, in mil_convert_to_proto
    prog = frontend_converter(model, **kwargs)
  File "/home/freedom/tf-master/lib/python3.8/site-packages/coremltools/converters/mil/converter.py", line 75, in __call__
    return tf2_loader.load()
  File "/home/freedom/tf-master/lib/python3.8/site-packages/coremltools/converters/mil/frontend/tensorflow/load.py", line 59, in load
    self._graph_def = self._graph_def_from_model(outputs)
  File "/home/freedom/tf-master/lib/python3.8/site-packages/coremltools/converters/mil/frontend/tensorflow2/load.py", line 99, in _graph_def_from_model
    saved_model = _tf.saved_model.load(self.model)
  File "/home/freedom/tf-master/lib/python3.8/site-packages/tensorflow/python/saved_model/load.py", line 603, in load
    return load_internal(export_dir, tags, options)
  File "/home/freedom/tf-master/lib/python3.8/site-packages/tensorflow/python/saved_model/load.py", line 649, in load_internal
    root = load_v1_in_v2.load(export_dir, tags)
  File "/home/freedom/tf-master/lib/python3.8/site-packages/tensorflow/python/saved_model/load_v1_in_v2.py", line 263, in load
    return loader.load(tags=tags)
  File "/home/freedom/tf-master/lib/python3.8/site-packages/tensorflow/python/saved_model/load_v1_in_v2.py", line 188, in load
    meta_graph_def = self.get_meta_graph_def_from_tags(tags)
  File "/home/freedom/tf-master/lib/python3.8/site-packages/tensorflow/python/saved_model/load_v1_in_v2.py", line 76, in get_meta_graph_def_from_tags
    raise ValueError(
ValueError: Importing a SavedModel with tf.saved_model.load requires a 'tags=' argument if there is more than one MetaGraph. Got 'tags=None', but there are 2 MetaGraphs in the SavedModel with tag sets [['serve'], ['serve', 'tpu']]. Pass a 'tags=' argument to load this SavedModel.

Nonetheless, I could not find a way to pass the tags all the way from ct.convert() to tf.saved_model.load()

I didn't try to add tags to the argument list because it's a tf saved_model specific one. If you think it's OK, I'll add it to the list and add a simple tf saved_model file for unit test.

[aseemw]

Yeah that would be great if you can add a test case (say in https://github.com/apple/coremltools/blob/master/coremltools/converters/mil/frontend/tensorflow2/test/test_v2_load.py or https://github.com/apple/coremltools/blob/master/coremltools/test/api/test_api_examples.py).

I think actually, adding it as part of kwargs initially might be better, before jumping in to add it as an official argument. We can document it in the doc string in _converters_entry.py. Later, it can be moved to the official argument list as well.