使用 Keras 的 ImageDataGenerator 划分训练集和测试集

Keras的ImageDataGenerator可以方便的读入文件夹中的图片并自动生成训练数据 . 如果图片已经被分成训练集和测试集两个文件夹, 可以独立创建两个ImageDataGenerator. 但是在很多情况下, 作者提供的数据集并不区分训练集和测试集, 这时候也可以使用ImageDataGenerator.

假设当前目录下有一个 Mushrooms数据集, 里面有9个子文件夹, 分别是9种蘑菇. 可以使用以下代码将生成ImageDataGenerator, 把数据集分成70%训练集和30%测试集.

注意代码中关键的validation_split和subset

data_dir = 'Mushrooms'
train_datagen = ImageDataGenerator(
        rescale=1./255,
        shear_range=0.2,
        zoom_range=0.2,
        validation_split = 0.3,
        horizontal_flip=True)
train_generator = train_datagen.flow_from_directory(
        data_dir,
        target_size=(150, 150),
        batch_size=32,
        class_mode='categorical',
        subset='training')

validation_generator = train_datagen.flow_from_directory(
    data_dir, # same directory as training data
    target_size=(150, 150),
    batch_size=32,
    class_mode='categorical',
    subset='validation') # set as validation data
    ```