tensorflow: TFLiteSavedModelConverterV2 breaks ordering of multi output models

@abattery Thanks for re-opening the issue. I tried out using SignatureDefs with Keras models and discovered a few bugs and usability issues which can be reproduced in this colab notebook.

The first problem is related to using full int8 quantization (/cc @MeghnaNatraj) by setting

converter.inference_input_type = tf.int8
converter.inference_output_type = tf.int8

This breaks the signature runner since the signature still refers to tensors that have been removed when converting to full int8 quantization.

The second problem is more a usability issue. Since the signature input and output names are set via Keras layer names, the signature of a Keras model depends on whether layers with the same name have already been instantiated prior to model saving which would lead to a change in leayer names. I think this is problematic since users will need to rely on fixed names in their TFLite code when relying on signature defs as far as I can tell. Otherwise this means during training great care needs to be taken that the model is saved in a fresh environment to not mess up the names (this can be very tricky during fine tuning where multiple instances of the same model might exist in the training code).

Let’s construct a model using dictionaries as inputs and outputs which would provide a unique mapping in multi-output models that does not rely on the ordering

model = tf.keras.Model({"input": input_tensor}, {"boxes": boxes, "scores": scores})

print(model.input)
# {'input': <KerasTensor: shape=(None, 32, 32, 128) dtype=float32 (created by layer 'input_4')>}

print(model.output)
# {'boxes': <KerasTensor: shape=(None, 4) dtype=float32 (created by layer 'tf.reshape_6')>,
# 'scores': <KerasTensor: shape=(None,) dtype=float32 (created by layer 'tf.reshape_7')>}

Unfortunately the signature def would include layer names that might not be uniquely identifiable across model training runs.

print(interpreter.get_signature_list()["serving_default"])
# {'inputs': ['input_4'], 'outputs': ['tf.reshape_6', 'tf.reshape_7']}

Please take a look at this colab notebook for a fully reproducible code example.

For me, all of these issues would be solved by maintaining the input/output order of the Keras model during serialization and conversion to TFLite so that the order of model.inputs and model.outputs (returns a list of tensors) are preserved during conversion. Alternatively, relying on SignatureDefs would work very well too, if it supports full int8 conversion and if there exists an easy way to repeatedly name inputs and outputs during Keras model saving that does not depend on the automatic layer naming of Keras.

For the multiple input/output naming/ordering issue, we recommend using TFLite signature def API.

SignatureDef provides meaningful/generic names for inputs/outputs which doesn’t rely on specific model details. More on SignatureDefs here. If your saved model has a defined signatureDef then it will be exported during conversion to TFLite and then you can use the Signature inputs/outputs for inference and not relying on inputs/outputs order or tensor names.

SignatureDef support is new thing available in nightly and will be available in 2.5

Then use the nightly for converting to tflite. The generated tflite file will have the SignatureDef details. Using Python API (as an example) you can do something like

my_signature = interpreter.get_signature_runner("my_method")
results = my_signature(input_1=input_tensor_1, input_2=input_tensor_2)
print(results["my_output"])

The signature def API is designed intentionally without having an intuitive order. It would be better to rely on the names instead of the order.

I agree that using the signature def API makes a lot of sense when working with tf.functions and saved models directly.

However, I think users only working with Keras models tend to expect that output ordering is preserved similar to how Python functions returning tuples work. So when converting models with tf.lite.TFLiteConverter.from_keras_model the fact that output ordering is not preserved is surprising since the fact that models get serialized as saved models during conversion is only an implementation detail that changed in TF 2.4. Is there a concrete reason for why preserving output ordering is not possible for Keras models? When saving and loading Keras models using model.save and tf.keras.models.load_model the ordering is preserved as well, which is why I am confused that this is not possible when converting to TFLite.

I am currently facing the same issue as lgeiger@

As shown in this comment above, once I train a keras model and specify input and output names, the input name stays as is but the output name is not preserved.

I am guessing the reason is because the outputs are tf.Tensor and not a tf.keras.layer.Layer type.

To overcome this, I have added a tf.keras.layer.Identity to each of my outputs.

eg: In the above scenario, that would translate to:

scores = tf.keras.layer.Identity(scores, name="scores)
boxes = tf.keras.layer.Identity(boxes, name="boxes)

model = tf.keras.Model({"input": input_tensor}, {"boxes": boxes, "scores": scores})

This ensures that the signature runner names are correct:

print(interpreter.get_signature_list()["serving_default"])
# {'inputs': ['input_4'], 'outputs': ['scores', 'boxes']}

However, as TFLite models are usually operating under smaller resource constraints, adding another layer doesn’t seem like a good solution. Ideally, it would be better if the issue is fixed in TFLite (ensure signature names are the same as the name of the output tensor)

Thanks @lgeiger for finding a bug! I took a look at the above colab. It looks like the signature API can be broken after the model optimization is applied. Could you copy and paste the comment (https://github.com/tensorflow/tensorflow/issues/47927#issuecomment-804240854) to create a separate issue? I will take a look at this issue with the team.

Actually, you can override the output signature names by wrapping the keras model.

For example,

@tf.function(input_signatures=[tf.TensorSpec([None, 32, 32, 128], dtype=tf.float32, name="input_tensor")])
def override_output_signatures(input_tensor):
   outputs = keras_model(input_tensor)
   return {"boxes": outputs.boxes, "scores": outputs.scores}

signatures = override_output_signatures.get_concrete_function()
tf.saved_model.save(model, export_dir=export_path, signatures=signatures)

For the above cases, I would recommend exporting your keras model to the corresponding saved model with a signature def before putting it to the TensorFlow Lite converter API for now to avoid the multi output names/ordering issue.

Thank you for providing the valuable feedback. At least, there should be improvements in the document side.