使用ImageDataGenerator时无需调整文件夹如何修改分类数量？

我现在用的自定义数据集包含4个类别（猫、狗、老鼠、金枪鱼），已经能很好完成这4类的分类任务。现在想把任务改成区分哺乳类与非哺乳类，但找不到不用修改图像存储文件夹结构的实现方法。

我目前想到的低效方法 #

我写了一个逐样本处理的生成器，但效率较低：

def to2Clases(image_generator,batch_size):
  noMammal= (0,0,0,1,)
  for x, y in image_generator:
    for i in range(len(y)):
      
      if all(y[i]==noMammal): y[i] = (0,1,0,0,)
      else: y[i] = (1,0,0,0,)
  yield x, y

注：原代码存在逻辑错误，yield 应放在外层循环内部，否则只会返回最后一个batch的数据。

我也考虑过先按原类别预测再转换结果，但本质和这个方法一样，效率不高。

参考过但不适用的方法 #

我找到一个基于MNIST数据集转奇偶二分类的方案，但它不适用于我的场景：

even=[0,2,4,6,8]

# Converting Train labels from numbers to even/odd
Y_trainbinary= np.empty(Y_train.shape[0],dtype=object)     
for idx, item in enumerate(Y_train):                     
      if np.nonzero(item)[0] in even:
            Y_trainbinary[idx] = 0     # even number
      else:
            Y_trainbinary[idx] = 1     # odd number

Y_trainmodel = np_utils.to_categorical(Y_trainbinary)

这个方案针对的是标签编码的场景，而我的训练生成器采用的是One-Hot编码，无法直接复用。

补充信息 #

我的训练生成器定义如下：

train_datagen = ImageDataGenerator(
    rescale = 1./255,
    horizontal_flip = True,
    vertical_flip = True,
    validation_split = 0.1)

train_generator = train_datagen.flow_from_directory(
    data_train,
    target_size = (img_height, img_width),
    batch_size = batch_size,
    shuffle = True,
    class_mode = 'categorical',
    subset = 'training',
    seed = 42)

打印标签形状的代码及输出：

x,y = train_generator.next()

print(len(y))
print(len(y[0]))
print(y[0])

输出：

32  # batch size
4   # number of classes
[0. 1. 0. 0.]  # 单样本One-Hot标签

更优的实现方法 #

可以利用Numpy的向量化操作替代逐样本循环，大幅提升效率，同时无需修改文件夹结构：

方法1：生成器层面转换标签（推荐） #

import numpy as np

def convert_to_mammal_binary(generator):
    # 定义原类别到哺乳类的映射：原类别索引 -> 1（哺乳类）/0（非哺乳类）
    # 需与`train_generator.class_indices`的顺序对应，假设顺序为猫、狗、老鼠、金枪鱼
    class_mapping = np.array([1, 1, 1, 0])
    
    for x, y in generator:
        # 获取每个样本的原类别索引（从One-Hot编码转换）
        original_class_indices = np.argmax(y, axis=1)
        # 映射为二分类标签
        binary_labels = class_mapping[original_class_indices]
        # 转换为One-Hot编码（若模型输出层为2分类的softmax）
        binary_one_hot = np.stack([1 - binary_labels, binary_labels], axis=1)
        # 若模型用sigmoid做二分类，直接返回binary_labels即可
        yield x, binary_one_hot

使用时直接包装原生成器：

train_binary_generator = convert_to_mammal_binary(train_generator)

方法2：修改模型输出层（适合微调场景） #

如果想复用原有模型的权重，可以在原4分类输出层后添加一个全连接层，将4维输出映射为2维：

from tensorflow.keras.models import Model
from tensorflow.keras.layers import Dense

# 假设原模型为`original_model`，输出层名为`dense_4`
base_model = Model(inputs=original_model.input, outputs=original_model.get_layer('dense_4').output)
# 添加二分类输出层
output_layer = Dense(2, activation='softmax')(base_model.output)
binary_model = Model(inputs=base_model.input, outputs=output_layer)

# 编译并训练
binary_model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])
binary_model.fit(train_generator, epochs=10)

这种方法无需修改生成器，直接在模型层面完成类别映射，适合基于原有模型进行微调的场景。

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