How can I convert "BigGAN (TensorFlow 1)" to TensorFlow Lite format?

I’ve also asked my question on stackoverflow, but I thought this forum might be better for Kaggle-related questions, so I’m posting here.

How can I convert “BigGAN (TensorFlow 1)” to TensorFlow Lite format?

I want to download “BigGAN (TensorFlow 1)” from Kaggle.Model and save it in TensorFlow Lite format. However, when I write the following source code, the following “TypeError” occurs.
How should I use KerasLayer with multiple inputs? I’m also not sure about the arguments to specify in “model.build”.

Non-working source code

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import tensorflow_hub as hub

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

### -->> TypeError occur HERE
output = hub_layer([input_truncation, input_y, input_z])
### <<-- TypeError occur HERE

model = tf.keras.models.Model(inputs=[input_truncation, input_y, input_z], outputs=[output])

### -->> How to build this model 
model.build([1], [1, 1000], [1, 128])
### <<-- How to build this model

model.summary()
model.save("biggan-128")

the exception that occurred

Exception encountered when calling layer "keras_layer" (type KerasLayer).

in user code:

    File "/home/shino/anaconda3/envs/movenet/lib/python3.10/site-packages/tensorflow_hub/keras_layer.py", line 242, in call  *
        result = f()

    TypeError: pruned(truncation, y, z) takes 0 positional arguments, got 1.


Call arguments received by layer "keras_layer" (type KerasLayer):
  • inputs=['tf.Tensor(shape=(None,), dtype=float32)', 'tf.Tensor(shape=(None, 1000), dtype=float32)', 'tf.Tensor(shape=(None, 128), dtype=float32)']
  • training=None
  File "/tmp/__autograph_generated_file8njdbc6p.py", line 74, in tf__call
    ag__.if_stmt(ag__.not_(ag__.ld(self)._has_training_argument), if_body_3, else_body_3, get_state_3, set_state_3, ('result', 'training'), 1)
  File "/tmp/__autograph_generated_file8njdbc6p.py", line 37, in if_body_3
    result = ag__.converted_call(ag__.ld(f), (), None, fscope)
TypeError: pruned(truncation, y, z) takes 0 positional arguments, got 1.

During handling of the above exception, another exception occurred:

  File "/tmp/__autograph_generated_file8njdbc6p.py", line 74, in tf__call
    ag__.if_stmt(ag__.not_(ag__.ld(self)._has_training_argument), if_body_3, else_body_3, get_state_3, set_state_3, ('result', 'training'), 1)
  File "/tmp/__autograph_generated_file8njdbc6p.py", line 37, in if_body_3
    result = ag__.converted_call(ag__.ld(f), (), None, fscope)
TypeError: in user code:

    File "/home/shino/anaconda3/envs/movenet/lib/python3.10/site-packages/tensorflow_hub/keras_layer.py", line 242, in call  *
        result = f()

    TypeError: pruned(truncation, y, z) takes 0 positional arguments, got 1.


During handling of the above exception, another exception occurred:

  File "/tmp/__autograph_generated_file8njdbc6p.py", line 37, in if_body_3
    result = ag__.converted_call(ag__.ld(f), (), None, fscope)
  File "/tmp/__autograph_generated_file8njdbc6p.py", line 74, in tf__call
    ag__.if_stmt(ag__.not_(ag__.ld(self)._has_training_argument), if_body_3, else_body_3, get_state_3, set_state_3, ('result', 'training'), 1)
  File "/home/shino/sandbox/python/biggan/biggan_export.py", line 37, in <module>
    output = hub_layer([input_truncation, input_y, input_z])
TypeError: Exception encountered when calling layer "keras_layer" (type KerasLayer).

in user code:

    File "/home/shino/anaconda3/envs/movenet/lib/python3.10/site-packages/tensorflow_hub/keras_layer.py", line 242, in call  *
        result = f()

    TypeError: pruned(truncation, y, z) takes 0 positional arguments, got 1.


Call arguments received by layer "keras_layer" (type KerasLayer):
  • inputs=['tf.Tensor(shape=(None,), dtype=float32)', 'tf.Tensor(shape=(None, 1000), dtype=float32)', 'tf.Tensor(shape=(None, 128), dtype=float32)']
  • training=None

By modifying the source code as follows, we were able to successfully convert BigGAN into TensorFlow Lite format.

Thank you.

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")

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