基于带二分类标签的数组创建ImageDataGenerator的正确方法

1 核心错误原因 #

先解释报错中的logits：指模型最终输出层的计算结果，你这里最后一层加了softmax激活，实际输出的是每个样本属于2个类别的概率值，TensorFlow报错中会统一将模型最终输出称为logits。

报错的根源有2个：

• 损失函数与标签格式不匹配：你当前的标签是one-hot编码的(样本数, 2)格式，而你选用的sparse_categorical_crossentropy损失仅支持输入一维整数标签（形状为(样本数,)，每个值为0或1）。当你把形状为(32,2)的批次标签喂给该损失时，函数会自动将标签展平为(64,)的一维数组，因此出现了logits形状[32,2]和标签形状[64]不匹配的问题。
• 输入形状写反：你加载的图片尺寸为(216, 172, 3)，但模型第一层Conv2D的input_shape参数写反为(172, 216, 3)，会导致后续计算的形状不匹配。

2 解决方法 #

两种方案二选一即可，不需要修改生成器代码，也不用换成flow_from_directory()，不影响后续添加图像增强逻辑：

方案一：保留当前one-hot标签，更换损失函数 #

直接修改模型compile阶段的损失函数为categorical_crossentropy，该损失原生支持one-hot格式的标签，同时修正输入形状：

model = tf.keras.Sequential([
    # 修正input_shape为实际图片尺寸
    tf.keras.layers.Conv2D(32, (3, 3), activation='relu', input_shape=(216, 172, 3)),
    tf.keras.layers.MaxPooling2D(2, 2),
    tf.keras.layers.Conv2D(64, (3, 3), activation='relu'),
    tf.keras.layers.MaxPooling2D(2, 2),
    tf.keras.layers.Conv2D(128, (3, 3), activation='relu'),
    tf.keras.layers.MaxPooling2D(2, 2),
    tf.keras.layers.Conv2D(128, (3, 3), activation='relu'),
    tf.keras.layers.MaxPooling2D(2, 2),
    tf.keras.layers.Dropout(0.5),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(512, activation='relu'),
    tf.keras.layers.Dense(2, activation='softmax')
])
model.compile(
    optimizer='adam',
    # 更换损失函数为适配one-hot标签的分类交叉熵
    loss='categorical_crossentropy',
    metrics=['accuracy']
)

修改完成后你可以直接在ImageDataGenerator中添加旋转、平移、翻转等图像增强参数，不需要调整其他逻辑。

方案二：转换标签为整数格式，保留原有损失函数 #

如果你想继续使用sparse_categorical_crossentropy，只需要在加载标签时将one-hot格式转换为一维整数格式即可：

# 标签转换代码，在拆分数据集前执行即可
labels = np.argmax(labels, axis=1)
# 后续拆分数据集、生成器代码均无需修改，模型input_shape同样要修正为(216, 172, 3)

3 关于class_mode参数的补充说明 #

你之前的猜测并不完全正确，flow()方法支持自动适配标签格式，不需要设置class_mode参数。flow_from_directory()的class_mode参数本质作用就是控制输出标签的格式：

• class_mode='binary'输出一维整数标签，适配binary_crossentropy或sparse_categorical_crossentropy
• class_mode='categorical'输出one-hot格式标签，适配categorical_crossentropy
  你用flow()方法直接传入自己构造的标签数组，生成器会直接输出你传入的标签格式，不需要额外配置。

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