python - 调用子层时，带有输入签名的 tf.function 出现错误

Question

以下效果很好

class MyModel(tf.Module):
    def __init__(self):
        super(MyModel, self).__init__(name='MyModel')
        self.dense = tf.keras.layers.Dense(50)

    @tf.function
    def __call__(self, inputs):
        return self.dense(inputs)

my_model = MyModel()
my_model(tf.ones((4, 42, 30, 200), name='features'))

但是，当我尝试在 tf.function 装饰器中显式指定输入签名时，代码失败:

class MyModel(tf.Module):
    def __init__(self):
        super(MyModel, self).__init__(name='MyModel')
        self.dense = tf.keras.layers.Dense(50)

    @tf.function(input_signature=[
        tf.TensorSpec(shape=(None, None, None, None), dtype=tf.float32)
    ])
    def __call__(self, inputs):
        return self.dense(inputs)

my_model = MyModel()
my_model(tf.ones((4, 42, 30, 200), name='features'))

导致以下错误:

ValueError: The last dimension of the inputs to `Dense` should be defined. Found `None`.

有人可以解释一下我在这里做错了什么吗？

Answer 1

您的tf.Module正在创建 tf.keras.Model ; Keras 模型有两个不同的阶段:

1. 模型构建，其中每一层都是根据输入形状定义的
2. 模型执行，其中构建的模型用于运行前向传递。

当您指定模型第一层的输入形状(使用 input_shape 或 Dense 层的 tf.keras.layers.Input 属性)时，Keras 能够执行第一阶段，即模型构建，因为它现在知道第一个权重矩阵应该类似于 input_shape[-1] * <number_units> (你的情况是 50)。

当您不指定输入形状时，会发生的情况是 Keras 无法在模型定义时间执行模型构建，但它已经等待第一个输入传递给模型；该输入具有已定义的形状(在您的情况下为 (4, 42, 30, 200) )，因此它可以连续运行模型构建和执行。

当您没有指定形状时，所有这些都会发生。

但在第二个示例中，您指定 __call__ 的输入形状你的方法tf.Module传播到 __call__您的tf.keras.Model .

因此，这里发生的情况是，您要求 Keras 模型连续进行模型构建和执行，但由于没有定义形状(因为您传递 (None, None, None, None) ，因此 input_shape[-1] is None ) Keras 只是无法定义模型。