如何将在自定义循环中训练的子类tf.keras.model转换为tflite？

Question

我在自定义循环中将经过训练的子类模型（tf.keras.Model）转换为TFLite时遇到问题。

假设我们有一个小的CNN架构，该架构使用输入数据（x）和其他信息，这些信息取决于批处理大小和其他尺寸（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

我以自定义循环模式（使用tf.GradientTape）对其进行了训练。这意味着我没有编译模型，而是按原样使用它。

现在，我想将其转换为SavedModel格式，因为我想在TFLite中移植我的模型。但是当我运行类似的内容时：

tf.saved_model.save(model, path_to_file)

我收到如下警告： Skipping full serialization of Keras model <ContextExtractor object at 0x7f30340bd6d8>, because its inputs are not defined.

当然，还有我得到的.pb文件-很小，里面什么也没有。

谁能提供完整的解释如何将子类模型转换为SavedModel？或者，也许我可以在没有它的情况下将模型转换为TFLite？