如何将在自定义循环子类 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？