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的调用替换为对其他tensorflow操作的调用可以很好地工作。

Answer 1

这似乎是因为当输入是keras对称输入时，tf.map_fn无法确定输入张量的TypeSpec。

现在，我已经多次遇到tensorflow操作和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)

Answer 2

为什么不使用tf.keras.layers.Lambda层呢？

import tensorflow as tf

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

x = Input(shape=(10,))

y = tf.keras.layers.Lambda(lambda x : x * 2)(x)

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