TensorFlow中如何压缩MapDataset维度以适配Keras模型输入

报错根因 #

你的Keras模型要求输入形状为(None, 24)，但MapDataset输出的输入、标签张量形状均为(None, 1, 24)，多了一个大小为1的第二维度，因此触发维度校验失败。

之前方法失效原因 #

• 你使用的切片[:, -1:, :]保留了第二维度的冒号，切片后形状仍然是(None, 1, 24)，没有实现降维
• 若直接使用tf.squeeze()不指定axis参数，默认会删除所有大小为1的维度，若某个批次的批量大小为1时会误删批量维度，导致维度错乱

正确的MapDataset降维方案 #

对数据集调用map方法，显式删除输入、标签张量中axis=1的维度即可：

import tensorflow as tf
from tensorflow import keras

# 假设你的原始MapDataset变量名为dataset
def remove_extra_dim(x, y):
    # 显式指定删除axis=1的大小为1的维度，不会影响其他维度
    x_processed = tf.squeeze(x, axis=1)
    y_processed = tf.squeeze(y, axis=1)
    return x_processed, y_processed

# 应用维度转换到整个数据集
dataset = dataset.map(remove_extra_dim)

维度验证 #

处理完成后可以打印数据集的元素规格确认维度正确：

print(dataset.element_spec)
# 预期输出：
# (TensorSpec(shape=(None, 24), dtype=tf.float32, name=None), TensorSpec(shape=(None, 24), dtype=tf.float32, name=None))

备选方案：不改数据集，直接适配模型 #

如果不想修改数据集处理逻辑，可以直接在模型输入层后加一个Squeeze层适配维度：

# 输入层适配数据集的原始维度
input_ = keras.layers.Input(shape=(1, 24))
# 删掉多余的第二维度
squeezed_input = keras.layers.Squeeze(axis=1)(input_)
# 后续层逻辑和你原有代码一致
hidden1 = keras.layers.Dense(30, activation="relu")(squeezed_input)
hidden2 = keras.layers.Dense(30, activation="relu")(hidden1)
concat = keras.layers.concatenate([squeezed_input, hidden2])
output = keras.layers.Dense(1)(concat)
model = keras.models.Model(inputs=[input_], outputs=[output])

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