tf keras upsample2d operator produces wrong results on tflite conversion

[anilsathyan7]

System information

• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 18.04
• TensorFlow installed from (source or binary): source and binary
• TensorFlow version (or github SHA if from source): 2.3.0, tf-nightly, 2.5.0(latest source)

Command used to run the converter or code if you’re using the Python API
If possible, please share a link to Colab/Jupyter/any notebook.
Colab Link: https://colab.research.google.com/drive/1QZo8QZubmiiXIxc0O8SZEfEsfaxN5kxR?usp=sharing

converter = tf.lite.TFLiteConverter.from_keras_model(model)
tflite_model = converter.convert()
open("mnv3_test.tflite", "wb").write(tflite_model)

Failure details
The conversion is successful, but the generated model is wrong,

The tf keras model with upsample2d layers(bilinear interpolation) correctly produces a h5 file; but the converted tflite model has extra layers and there is shape mismatch.The output shape becomes 1x1 instead of original upsampled size. The error occurs in tf 2.3.0, tf nightly and latesr source (dev 2.5.0). The conversion produces additonal shape, mul and strided slice layer(not supported by hardware accelerators).

The problem persists even if we use ResizeNearestNeighbor interpolation.

Minimal Example

import tensorflow as tf
from tensorflow.keras.layers import Input, Conv2D, UpSampling2D
from tensorflow.keras.models import Model
from tensorflow.keras.optimizers import Adam

def get_model():
    
    ip = Input(shape=(256,256,3))
    x= Conv2D(kernel_size=3, filters=32, padding ='same')(ip)
    x = UpSampling2D( size=(2, 2), interpolation='bilinear')(x)
    x= Conv2D(kernel_size=3, filters=32, padding ='same')(x)
    
    
    model = Model(inputs=ip, outputs=x)
    
    model.compile(optimizer=Adam(lr=1e-4),
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy']) # Ensure you have sparse labels
    return model

model=get_model()
converter = tf.lite.TFLiteConverter.from_keras_model(model)
tflite_model = converter.convert()
open("bilinear.tflite", "wb").write(tflite_model)

Models
mnv3_resize_bilinear_test.zip

[Saduf2019]

@ymodak
I am able to replicate the issue on nightly, please find the gist here.Thanks!

[sunzhe09]

I met the same problem, additonal shape, mul and strided slice layer(not supported by hardware accelerators).

[sunzhe09]

@Saduf2019 when I set my batch dim with 1 these layers was disappear;

[anilsathyan7]

@sunzhe09 How did you set batch dim to 1 ? The tfllite already has batch size 1 right?
Also, is the upsampled size correct this time?

[anilsathyan7]

@Saduf2019 If we set the batch dimesion to input layer during model defintion of tf.keras model, all the problems seems to be resolved. However can we set the batch dimension of all layers for the already trained model to 1, before tflite conversion?

[ymodak]

The minimal example generated tflite file viewed in netron shows additonal shape, mul and strided slice layer.

[karimnosseir]

This is similar to issue #43882. See my explanation on why this is correct and an example on how to set your model shape to be static before conversion

#43882 (comment)

[Vozf]

For anyone wondering. For single input keras model the fix above can be implemented like this.

model.input.set_shape(1 + model.input.shape[1:])