How can I create a model in tflite format with a fixed batch_size?

TensorFlow General Discussion
, , ,
1

Hello.

I am trying to create a tflite format model from BigGAN downloaded from Kaggle Hub. When creating the model, I want to fix the input_shape (batch_size) to 1 due to the specifications of the APU of the embedded device. I wrote the following code, but it does not process KerasLayer properly. How can I convert a model that includes KerasLayer to tflite format with fixed input_shape?

Source code where the problem occurs

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import gridspec
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import tensorflow_hub as hub
import tf2onnx
import onnx
import onnxruntime
from onnxruntime.datasets import get_example

model_url = "https://www.kaggle.com/models/deepmind/biggan/TensorFlow1/128/2"
input_truncation = keras.Input(batch_shape=(), name='truncation')
input_y = keras.Input(shape=(1000, ), name='y')
input_z = keras.Input(shape=(120, ), name='z')
hub_layer = hub.KerasLayer(
            model_url,
            trainable=False,
            signature="default", 
            signature_outputs_as_dict=True,
        )

inputs = dict(y=input_y, z=input_z, truncation=input_truncation)
output = hub_layer(inputs)
model = tf.keras.models.Model(inputs=[input_truncation, input_y, input_z], outputs=[output])
model.build(([1], [1, 1000], [1, 120]))
model.summary()
model.save("biggan-128")
model.save("biggan-128.h5")

### -->> ERROR
converter = tf.compat.v1.lite.TFLiteConverter.from_keras_model_file("biggan-128.h5", input_shapes={'truncation[0][0]':(), 'y[0][0]':(1, 1000), 'z[0][0]':(1, 120)})
### <<-- ERROR
tflite_model = converter.convert()
open("biggan-128.tflite", "wb").write(tflite_model)

Error message

ValueError
Unknown layer: KerasLayer. Please ensure this object is passed to the `custom_objects` argument. See https://www.tensorflow.org/guide/keras/save_and_serialize#registering_the_custom_object for details.
  File "/home/shino/sandbox/python/biggan/biggan_export.py", line 32, in <module>
    converter = tf.compat.v1.lite.TFLiteConverter.from_keras_model_file("biggan-128.h5", input_shapes={'truncation[0][0]':(), 'y[0][0]':(1, 1000), 'z[0][0]':(1, 120)})
ValueError: Unknown layer: KerasLayer. Please ensure this object is passed to the `custom_objects` argument. See https://www.tensorflow.org/guide/keras/save_and_serialize#registering_the_custom_object for details.

What has been confirmed

By the way, when I created the following code, I was able to create a tflite format model with dynamic input_shape.

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

requirements.txt

absl-py==2.2.1
astunparse==1.6.3
cachetools==5.5.2
certifi==2025.1.31
charset-normalizer==3.4.1
coloredlogs==15.0.1
contourpy==1.3.1
cycler==0.12.1
flatbuffers==1.12
fonttools==4.57.0
gast==0.4.0
google-auth==2.38.0
google-auth-oauthlib==0.4.6
google-pasta==0.2.0
grpcio==1.71.0
h5py==3.13.0
humanfriendly==10.0
idna==3.10
keras==2.9.0
Keras-Preprocessing==1.1.2
kiwisolver==1.4.8
libclang==18.1.1
Markdown==3.7
MarkupSafe==3.0.2
matplotlib==3.10.1
mpmath==1.3.0
numpy==1.23.5
oauthlib==3.2.2
onnx==1.14.1
onnx-graphsurgeon==0.5.7
onnx2tf==1.26.9
onnxruntime==1.21.0
opt_einsum==3.4.0
packaging==24.2
pillow==11.1.0
protobuf==3.20.3
psutil==7.0.0
pyasn1==0.6.1
pyasn1_modules==0.4.2
pyparsing==3.2.3
python-dateutil==2.9.0.post0
requests==2.32.3
requests-oauthlib==2.0.0
rsa==4.9
six==1.17.0
sng4onnx==1.0.4
sympy==1.13.3
tensorboard==2.9.0
tensorboard-data-server==0.6.1
tensorboard-plugin-wit==1.8.1
tensorflow==2.9.0
tensorflow-estimator==2.9.0
tensorflow-hub==0.16.1
tensorflow-io-gcs-filesystem==0.37.1
tensorflow-neuron==1.0
termcolor==3.0.0
tf-keras==2.14.1
tf2onnx==1.13.0
typing_extensions==4.13.0
urllib3==2.3.0
Werkzeug==3.1.3
wrapt==1.17.2
2

By using batch_shape when defining keras.Input,
I was able to create a model with a fixed batch size of 1.

Thank you.

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import gridspec
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import tensorflow_hub as hub
import tf2onnx
import onnx
import onnxruntime
from onnxruntime.datasets import get_example

model_url = "https://www.kaggle.com/models/deepmind/biggan/TensorFlow1/128/2"
input_truncation = keras.Input(batch_shape=(), name='truncation')
input_y = keras.Input(batch_shape=(1, 1000, ), name='y')
input_z = keras.Input(batch_shape=(1, 120, ), name='z')
hub_layer = hub.KerasLayer(
            model_url,
            trainable=False,
            signature="default", 
            signature_outputs_as_dict=True,
        )

inputs = dict(y=input_y, z=input_z, truncation=input_truncation)
output = hub_layer(inputs)
model = tf.keras.models.Model(inputs=[input_truncation, input_y, input_z], outputs=[output])
model.build(([1], [1, 1000], [1, 120]))
model.summary()
model.save("biggan-128")
model.save("biggan-128.h5")

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

onnx_model, _ = tf2onnx.convert.from_keras(model, outputs_as_nchw=[model.outputs[0].name])
onnx.save(onnx_model, './biggan-128.onnx')