如何设置tf.data.Dataset维度以适配带Masking层的LSTM时序分类任务

问题根因 #

你遇到的维度不匹配报错，核心是tf.data.Dataset的构造接口使用错误：

• tf.data.Dataset.from_tensors会把传入的完整数组作为单个样本处理，你传入的tf_data维度是(225970, 180, 1)，所以构造出的数据集每个元素的维度就是(225970, 180, 1)，再调用batch(64)后，输入维度就变成了报错里的(None, 225970, 180, 1)，比模型要求的3维输入多了一个维度，自然校验失败。

解决方法 #

只需要调整数据集的构造方式即可，不需要修改原有可运行的模型配置：

1. 把from_tensors替换为from_tensor_slices，这个接口会沿着输入数组的第0维切分，每个元素对应一条独立样本，刚好匹配你225970条样本的结构。
2. 不需要给模型输入层强行指定batch_size参数，tf.data的batch操作会自动处理batch维度，直接使用你原有的numpy版本模型结构即可。

可运行代码示例 #

数据集构造 #

# 按样本维度切分，每个元素为(单条样本, 对应标签)
tf_dataset = tf.data.Dataset.from_tensor_slices((tf_data, labels.astype(int)))
# 打乱、分批，追加prefetch可优化数据加载性能
tf_dataset = tf_dataset.shuffle(buffer_size=1024).batch(64).prefetch(tf.data.AUTOTUNE)

模型训练 #

直接复用你原有可正常运行的模型代码即可：

model = keras.Sequential(
    [
         layers.Masking(mask_value=0, input_shape=(180,1)),
         layers.LSTM(16),
         layers.Dense(1, activation='sigmoid')
    ]
)

model.compile(
    optimizer='adam',
    loss=tf.keras.losses.BinaryCrossentropy(from_logits=True),
    metrics=['accuracy']
)

model.fit(tf_dataset, epochs=10)

补充说明 #

如果你确实有需要在Input层指定batch_size的场景（比如使用状态化LSTM），要保证模型指定的batch_size和数据集的batch大小完全一致，普通训练场景不需要额外指定该参数。

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