在目录中的ImageDataGenerator流中将两个类视为一个

Question

我目前正在使用TensorFlow提供的数据生成器，我想通过将我的多类分类问题转换为二进制问题来做一些测试。

我的当前文件夹布局如下:

data/
├── 0/  <-- folder containing images of class 0
├── 1/
└── 2/

我目前正在加载它如下:

train_datagen = ImageDataGenerator(rescale=1./255, validation_split=0.5)
train_generator = train_datagen.flow_from_directory(
train_data_dir,
target_size=(input_size, input_size),
batch_size=batch_size,
class_mode='categorical',
subset='training')

当然我也有一个验证子集。但是，我正在努力弄清楚如何才能使它成为一个类，这样我就可以在不操纵实际数据或文件夹结构的情况下将类1和类2视为单个类。

我找到了一种方法来选择具体使用哪些类，但那是在文档中。

Answer 1

希望以下内容对您有所帮助:

train_datagen = ImageDataGenerator(rescale=1./255, validation_split=0.5)
train_generator = train_datagen.flow_from_directory(
train_data_dir,
target_size=(input_size, input_size),
batch_size=batch_size,
class_mode='categorical',
subset='training')
def transform_classes(args):
"""
data/0/ -> 0
data/1/ -> 1
data/2/ -> 1
"""
new_classes = np.array([0,1,1])
x,y = args
y = np.argmax(y, axis=1) # One hot to integer
y = new_classes[y] # Changing the target
y = np.eye(2)[y] # Going back to one hot 
return x,y
train_generator = map(transform_classes, train_generator)

这个答案背后的想法是train_datagen.flow_from_directory将返回一个python生成器，我们可以使用map进行转换。