TypeError：使用 tf.map_fn 和 keras 功能模型时无法为 <KerasTensor 构建 TypeSpec

Question

当我尝试在 keras 功能模型的定义中使用 tf.map_fn 时，出现错误：

TypeError: Could not build a TypeSpec for <KerasTensor: ...

例如，这个简单的模型将在 tf-nightly 2.5.0 中触发该错误：

import tensorflow as tf

from tensorflow.keras.layers import Input
from tensorflow.keras.models import Model

x = Input(shape=(10,))

y = tf.map_fn(lambda x : x * 2, x, fn_output_signature=tf.float32)

model = Model(inputs=x, outputs=y)

而将对 tf.map_fn 的调用替换为对其他张量流操作的调用则可以正常工作。

Answer 1

这似乎是由于当输入是 keras sympolic 输入时 tf.map_fn 无法确定输入张量的 TypeSpec。

现在我已经多次遇到张量流操作和 Keras 符号输入张量的问题。将有问题的代码包装在自定义层中似乎通常可以修复它。

例如，将上面的代码替换为以下内容即可成功执行：

import tensorflow as tf

from tensorflow.keras.layers import Input, Layer
from tensorflow.keras.models import Model

x = Input(shape=(10,))

class MapLayer(Layer):

    def call(self, input):
        return tf.map_fn(lambda x : x * 2, input, fn_output_signature=tf.float32)

y = MapLayer()(x)

model = Model(inputs=x, outputs=y)