`STFT` layer output shape deviates from `STFTTflite` layer in batch dimension

Use Case

I want to convert a STFT layer in my model to a STFTTflite to deploy it to my mobile device. In the documentation I found that another dimension is added to account for complex numbers. But I also encountered a behaviour that is not documented.

Expected Behaviour

input_shape = (2048, 1)  # mono signal

model = keras.models.Sequential()  # TFLite incompatible model
model.add(kapre.STFT(n_fft=1024, hop_length=512, input_shape=input_shape))

tflite_model = keras.models.Sequential()  # TFLite compatible model
tflite_model.add(kapre.STFTTflite(n_fft=1024, hop_length=512, input_shape=input_shape))

model has the output shape (None, 3, 513, 1). Therefore, tflite_model should have the output shape (None, 3, 513, 1, 2).

Observed Behaviour

The output shape of tflite_model is (1, 3, 513, 1, 2) instead of (None, 3, 513, 1, 2).

Problem Solution

• If this behaviour is unwanted:
  • Change the model output format so that the batch dimension is correctly shaped.
• Otherwise:
  • Explain in the documentation why the batch dimension is shaped to 1.
  • Explain in the documentation how to include this layer into models which expect the batch dimension to be shaped None.

[kenders2000]

Hi,

I can see in the model summary that while the input is None theSTFTTflite layers still have a batch size of 1. So while I would expect this model to convert and run fine (when you provide a batch size of one), you would still need resize the input dimension of the resulting tflite file to have a batch size of 1.

E.g. using resize_tensor_input()

Cheers

Paul

[keunwoochoi]

Thank you so much for everyone!