将NumPy数组转为tf.data.Dataset后Keras模型训练报错排查

问题代码 #

import numpy as np
import gc
import tensorflow as tf
from tensorflow import keras
import os

values = np.array([[0,0], [0,1], [1,0], [1,1]])
results = np.array([[0], [1], [1], [0]])

dataset = tf.data.Dataset.from_tensor_slices((values, results))
dataset

>>> <TensorSliceDataset shapes: ((2,), (1,)), types: (tf.int32, tf.int32)>

# Build a simple model
model = keras.Sequential()
model.add(keras.Input(shape=(2,)))
model.add(keras.layers.Dense(64, activation='relu')) # , kernel_initializer='he_uniform'
model.add(keras.layers.Dense(32, activation='relu'))
model.add(keras.layers.Dense(1, activation='sigmoid'))
model.compile(loss='binary_crossentropy', optimizer='adam', metrics=['accuracy'])
model.summary()

print("Fit on dataset/NumPy data")
history = model.fit(dataset, epochs=10)

>>> ...
>>> ValueError: Input 0 of layer sequential is incompatible with the layer: expected axis -1 of input shape to have value 2 but received input with shape (2, 1)

错误原因 #

直接用NumPy数组训练时，model.fit会自动将数组封装为批次数据（形状为(batch_size, 2)），与模型输入要求匹配。但tf.data.Dataset.from_tensor_slices创建的数据集默认是逐个样本输出，每个输入样本的形状是(2,)。

当你把这种未批量化的数据集传给model.fit时，TensorFlow会把单个样本的(2,)形状误判为「批次大小为2、每个样本维度为1」的输入（即形状(2, 1)），这与模型输入层要求的最后一维为2矛盾，因此抛出形状不兼容错误。

解决方法 #

给数据集添加批量化操作，让每个输出元素是一个批次的数据，形状与模型输入匹配：

# 创建数据集后添加batch操作，batch大小可按需调整
dataset = tf.data.Dataset.from_tensor_slices((values, results)).batch(4)
# 也可用更小的batch size，比如：
# dataset = tf.data.Dataset.from_tensor_slices((values, results)).batch(1)

# 执行训练
history = model.fit(dataset, epochs=10)

批量化后，数据集每个元素的输入形状为(4, 2)（或(1, 2)），完全符合模型输入层(None, 2)的要求（None表示支持任意批次大小），即可正常训练。

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