使用tensorflow_datasets训练鸢尾花分类模型遇形状不兼容错误

问题出在tensorflow_datasets（tfds）加载的数据集格式上：tfds默认返回的是单个样本组成的流式数据集，每个样本的输入是形状为(4,)的一维张量；而Keras模型训练时期望接收批量样本，要求输入形状为(batch_size, 4)的二维张量。这就是你看到"期望输入维度为2，实际得到维度1"错误的原因。

而用UCI CSV文件加载时，数据会被一次性读入成二维数组（比如(150,4)），model.fit会自动划分批次，因此不会出现形状不兼容的问题。

具体解决步骤 #

只需要对tfds加载的数据集添加批量处理，同时补充必要的打乱、重复操作，适配流式数据集的训练逻辑：

1. 预处理数据集

import tensorflow as tf
import tensorflow_datasets as tfds

# 加载数据集
ds = tfds.load('iris', split='train', shuffle_files=True, as_supervised=True)

# 打乱数据（缓冲区设为数据集总大小150，确保充分打乱）
# 批量处理（添加批量维度，匹配模型输入要求）
# 重复数据集（让每个epoch都能重新遍历数据）
ds = ds.shuffle(150).batch(50).repeat()

2. 训练模型（以Sequential API为例）

# 构建模型
model = tf.keras.Sequential([
    tf.keras.layers.Dense(16, activation='relu', input_shape=(4,)),
    tf.keras.layers.Dense(3, activation='softmax')
])

# 编译模型
model.compile(
    optimizer='adam',
    loss='sparse_categorical_crossentropy',
    metrics=['accuracy']
)

# 训练：steps_per_epoch设为总样本数除以batch_size（150//50=3）
model.fit(ds, steps_per_epoch=3, epochs=100)

Functional API示例

# 构建Functional模型
inputs = tf.keras.Input(shape=(4,))
x = tf.keras.layers.Dense(16, activation='relu')(inputs)
outputs = tf.keras.layers.Dense(3, activation='softmax')(x)
model = tf.keras.Model(inputs=inputs, outputs=outputs)

# 编译训练（逻辑和Sequential一致）
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
model.fit(ds, steps_per_epoch=3, epochs=100)

关键说明 #

• batch(50)：给每个样本添加批量维度，让输入形状从(4,)变为(50,4)，完全匹配模型输入层input_shape=(4,)的要求（模型会自动兼容批量维度）。
• shuffle(150)：确保每个epoch的样本顺序被打乱，避免模型过拟合。
• repeat()：让数据集在每个epoch结束后自动重置，保证model.fit能运行指定的100个epoch。
• steps_per_epoch：流式数据集需要明确每个epoch要执行多少步（总样本数/批量大小），避免模型无限遍历数据。

内容的提问来源于stack exchange，提问作者Vetle Hofsøy-Woie