python - 切片压缩的 tf.data.Dataset

Question

我有两个特征：名字和姓氏。我压缩它们以用作我的模型的多个输入。然后我批处理数据集。

import tensorflow as tf

names = tf.data.Dataset.from_tensor_slices(tf.ones(shape=(3,5,5)))
surnames = tf.data.Dataset.from_tensor_slices(tf.zeros(shape=(3,5,5)))
features = tf.data.Dataset.zip((names,surnames)).batch(3)

所以我得到的形状：

<BatchDataset shapes: ((None, 5, 5), (None, 5, 5)), types: (tf.float32, tf.float32)>

然后在通话中我想将它们分配给输入：

input_a = features.flat_map(lambda feature1,feature2:tf.data.Dataset.from_tensor_slices(feature1))

但是，如果我这样做，我会放弃批处理的概念，因为输入的形状变为 (5,5)
我该如何解决这个问题？