TensorFlow使用text_dataset_from_directory后如何将(10,1)调整为(10,)维度

问题原因 #

你加载数据集时先设置了batch_size=1，每个样本会被包装为(1,)形状的张量，后续调用.batch(10)拼接10个样本后，就会产生冗余的最后一维，得到(10,1)的形状，不符合模型输入要求。

解决方案 #

方案1：加载数据集时直接设置正确的批次大小（最推荐） #

直接在text_dataset_from_directory中指定所需的批次大小，无需后续二次分批，自然生成符合要求的张量形状：

train_data = tf.keras.preprocessing.text_dataset_from_directory(
    "data/text",
    batch_size=10 # 直接设置批次为10
)

train_features_batch, train_labels_batch = next(iter(train_data))
# 此时输出为 (10,) (10,)，符合模型输入要求
print(train_features_batch.shape, train_labels_batch.shape)

方案2：手动挤压冗余维度 #

如果需要保留原有分批逻辑，可以用tf.squeeze去掉大小为1的冗余维度：

train_features_batch, train_labels_batch = next(iter(train_data.batch(10)))
# 去掉最后一维的冗余维度
train_features_batch = tf.squeeze(train_features_batch, axis=-1)
train_labels_batch = tf.squeeze(train_labels_batch, axis=-1)

方案3：给数据集加全局预处理映射 #

如果需要对整个数据集的所有批次统一调整形状，可以通过map方法添加预处理逻辑，后续取数时自动适配：

# 全局处理所有样本的形状
train_data = train_data.map(lambda x, y: (tf.squeeze(x, axis=-1), tf.squeeze(y, axis=-1)))

train_features_batch, train_labels_batch = next(iter(train_data.batch(10)))
# 输出形状自动为 (10,) (10,)
print(train_features_batch.shape, train_labels_batch.shape)

以上方案都不需要修改你现有的模型结构，调整后的张量可以直接输入模型训练。

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