TensorFlow V2 to CoreML conversion fails

I'm trying to convert a TensorFlow model that I didn't create and know approximately nothing about to CoreML so that I can use it in some functional tests. I can't tell you much about the model, but you can read about it on the blog from the team that created it: https://research.google/blog/improving-mobile-app-accessibility-with-icon-detection/

I can't convert this model to a TensorFlow Lite model because it uses a few full TensorFlow operations (which I could work around) and it exceeds the 4-tensor output limit (which I can't, AFAIK). So instead, I'm trying to convert the model to CoreML so that I can run it on-device.

The issue I'm running into is that every approach fails in different ways. If I load the model with tf.saved_model.load and pass that as the first parameter to the convert call, it says

NotImplementedError: Expected model format: [SavedModel | concrete_function | tf.keras.Model | .h5 | GraphDef], got <tensorflow.python.trackable.autotrackable.AutoTrackable object at 0x30d90c250>

If I pass model.signatures['serving_default'] as the first parameter to convert, I get

NotImplementedError: Expected model format: [SavedModel | concrete_function | tf.keras.Model | .h5 | GraphDef], got ConcreteFunction [...a page or two of info about the function here...]

If I try to wrap it in a Keras layer using the instructions provided in the converter, it fails because a sequential model can't have multiple outputs.

If I try to use a tf.keras.layers.TFSMLayer to load the model, it fails because there are multiple tags, and there's no way to specify tags when constructing the layer. (It tells me that I need to add 'tags' to load the model, but if I do that, it tells me that tags isn't a valid parameter to the call.)

If I load the model with tf.saved_model.load and specify a single tag, then re-save it in a different location with tf.saved_model.save to generate a new model with only a single tag, then do

  input_layer = tf.keras.Input(shape=(768, 768, 3), dtype="int8")
  layer = tf.keras.layers.TFSMLayer("./serve_model", call_endpoint='serving_default')
  outputs = layer(input_layer)
  model = tf.keras.Model(input_layer, outputs)

I get

AttributeError: 'Functional' object has no attribute '_get_save_spec'

At one point, I also tried this:

  class LayerFromSavedModel(tf.keras.layers.Layer):
    def __init__(self):
      super(LayerFromSavedModel, self).__init__()
      self.vars = legacy_model.variables
    def call(self, inputs):
      return legacy_model.signatures['serving_default'](inputs)

  input = tf.keras.Input(shape=(3000, 3000, 3))
  model = tf.keras.Model(input, LayerFromSavedModel()(input))

and saw a similar failure.

I've run out of ideas here. Is there simply no support whatsoever in the converter for importing a TensorFlow 2 SavedModel into CoreML, or am I missing something fundamental?

Downgrading TensorFlow to 2.15 or 2.13 (the oldest available on Sonoma) got farther, but still fails with

tensorflow.python.framework.errors_impl.NotFoundError: Unable to find FunctionDef with name: __inference_cond_2_false_257083_439_tfg_region_specialized_map_1_while_non_max_suppression_padded_PartitionedCall_cond_2_1_tfg_region_specialized_map_1_while_non_max_suppression_padded_PartitionedCall_cond_2_1

So it looks like now maybe I'm running into missing TensorFlow operations that TFLite doesn't support. Any ideas?

TensorFlow V2 to CoreML conversion fails
 
 
Q