如何在TensorFlow中为Keras子类化模型指定输入？

嘿，我懂你想要的效果——既要用子类化模型的灵活性，又能像函数式API那样明确定义输入层对吧？其实有个很简单的办法，咱们可以把子类化模型和函数式API结合起来用，具体看下面的示例：

方法一：用输入层包装子类化模型 #

这是最灵活的方式，子类模型本身不需要做太多修改，只需要在外部创建输入层并关联模型即可：

import tensorflow as tf
from tensorflow.keras import Model

# 你的子类化模型保持原有逻辑
class SubModel(Model):
    def call(self, inputs):
        print(inputs)
        # 这里可以添加你的实际模型逻辑，比如全连接层、卷积层等
        return inputs  # 根据你的任务替换成真实输出

# 手动定义输入层，和函数式API一样
inputs = tf.keras.Input(shape=(100,))
# 实例化你的子类模型
sub_model_instance = SubModel()
# 用输入层调用子类模型，得到输出张量
outputs = sub_model_instance(inputs)
# 包装成完整的函数式风格模型
model = tf.keras.Model(inputs=inputs, outputs=outputs)

# 接下来就可以正常编译、训练了
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')
model.fit(data, labels, epochs=10, batch_size=32)

方法二：在子类模型内部定义输入层 #

如果你希望输入层和子类模型绑定在一起，可以在__init__方法里初始化输入层：

class SubModel(Model):
    def __init__(self, input_shape):
        super().__init__()
        # 在初始化时创建输入层
        self.custom_input = tf.keras.Input(shape=input_shape)
        # 这里添加你的其他层，比如：
        # self.dense_layer = tf.keras.layers.Dense(64, activation='relu')
    
    def call(self, inputs):
        print(inputs)
        # x = self.dense_layer(inputs)
        return inputs

# 实例化时传入输入形状
sub_model = SubModel(input_shape=(100,))
# 关联输入和输出，构建完整模型
outputs = sub_model(sub_model.custom_input)
model = tf.keras.Model(inputs=sub_model.custom_input, outputs=outputs)

为什么要这么做？ #

这样做的好处是，你既可以享受子类化模型带来的自定义逻辑灵活性（比如复杂的循环、条件分支），又能像函数式API那样明确控制输入形状，避免模型在fit时自动推断可能出现的问题（比如多输入、动态形状场景）。

内容的提问来源于stack exchange，提问作者Alexander