How to convert trained in custom loop subclassed tf.keras.model to tflite?

Question

I have a problem with converting trained subclassed model (tf.keras.Model) in custom loop to TFLite.

Suppose we have small CNN architecture that use input data (x) and additional information that depends on batch size and other dims (add_info):

class ContextExtractor(tf.keras.Model):
    def __init__(self):
        super().__init__()
        self.model = self.__get_model()

    def call(self, x, training=False, **kwargs):
        b, h, w, c = x.shape
        add_info = tf.zeros((b, h, w, c), dtype=tf.float32) 
        features = self.model(tf.concat([x, add_info], axis=-1), training=training)
        return features

    def __get_model(self):
        return self.__get_small_cnn()

    def __get_small_cnn(self):
        model = tf.keras.Sequential()
        model.add(layers.Conv2D(32, (3, 3), strides=(2, 2), padding='same'))
        model.add(layers.LeakyReLU(alpha=0.2))

        model.add(layers.Conv2D(32, (3, 3), strides=(2, 2), padding='same'))
        model.add(layers.LeakyReLU(alpha=0.2))

        model.add(layers.Conv2D(64, (3, 3), strides=(2, 2), padding='same'))
        model.add(layers.LeakyReLU(alpha=0.2))

        model.add(layers.Conv2D(128, (3, 3), strides=(2, 2), padding='same'))
        model.add(layers.LeakyReLU(alpha=0.2))

        model.add(layers.Conv2D(256, (3, 3), strides=(2, 2), padding='same'))
        model.add(layers.LeakyReLU(alpha=0.2))

        model.add(layers.GlobalAveragePooling2D())

        return model

I trained it in custom loop mode (using tf.GradientTape). It means that I didn't compile model, I just use it as it is.

Now I want to convert it to SavedModel format, because I want to port my model in TFLite. But when I run something like:

tf.saved_model.save(model, path_to_file)

I got warning like: Skipping full serialization of Keras model , because its inputs are not defined.

And, of course, my .pb file that I got -- super small and nothing in it.

Can anyone provide full explanation how to convert subclassed model to SavedModel? Or maybe I can convert my model to TFLite without it?

How to convert trained in custom loop subclassed tf.keras.model to tflite?

0 Answers