keras批量图片生成器generator：解决训练集过大无法一次喂入

 

参考链接：https://blog.csdn.net/mieleizhi0522/article/details/82191331

还记得在跑mnist的时候，训练集里有60000张黑白图片，每张图片为28*28，把所有图片塞到一个矩阵里面一点问题都没有。

现在就不一样了，我的数据集里大概有90000张彩色图片，大小就是384*512*3，这远远大于mnist的训练集，如果同时将它们塞进一个矩阵里，那么内存估计够呛。

所以，我们要学会利用keras里面的批量图片生成器generator，一批一批的读图片放进去训练。下面我写一个函数，把图片批量读取和模型的训练都放在这个函数里面。

from keras.preprocessing.image import ImageDataGenerator
# 喂入数据训练
#函数参数（定义的模型，训练集存放文件夹路径，测试集存放文件夹路径，
#batch_size大小，跑多少个epoch，训练集有多少张图片，模型保存路径）
def model_train(model, destination_train_dir, destination_test_dir, 
batch_size_set, epoch_set, train_all_counter,model_save_name):
    # 构建生成器，从文件夹中循环读取训练集和测试集
    #训练集生成器，参数rescale将图片像素归一化，结合sigmoid激活函数效果更好
    train_datagen = ImageDataGenerator(rescale=1. / 255)
    #测试集生成器
    test_datagen = ImageDataGenerator(rescale=1. / 255)
    #训练集大小/batch大小=每个epoch要读多少次数据
    steps_per_epoch_set = math.ceil(train_all_counter / float(batch_size_set))
    #测试集大小/batch大小
    validation_steps_set = math.ceil((test_all_counter / float(batch_size_set)))

    #生成训练batch（一般放四个参数：训练集文件夹路径，每张图片的size，每个batch的大小，以及标签的类型）
    train_generator = train_datagen.flow_from_directory(
        destination_train_dir,
        target_size=(384, 512),
        batch_size=batch_size_set,
        class_mode='binary'
    )
    # 生成测试batch
    validation_generator = test_datagen.flow_from_directory(
        destination_test_dir,
        target_size=(384, 512),
        batch_size=batch_size_set,
        class_mode='binary'
    )
    # ---------------------------------------------------------------------------------------------------------------------------------
    # 训练（五个参数：训练集生成器，训练每个epoch要抽取多少次图片，多少个epoch，测试集生成器，测试要抽取多少次数据）
    history = model.fit_generator(
        train_generator,
        steps_per_epoch=steps_per_epoch_set,
        epochs=epoch_set,
        validation_data=validation_generator,
        validation_steps=validation_steps_set
    )
    # ---------------------------------------------------------------------------------------------------------------------------------
    # 保存模型
    model_path = os.path.join(model_dir, model_save_name)
    model.save(model_path)
    return history

我们重点关注 train_generator = train_datagen.flow_from_directory（…）生成训练集batch的这一段，像代码中那样依次把参数都填进去就可以了生成相应训练batch。

需要强调的是，你要解决n分类问题，就要在destination_train_dir这个文件夹（你自己设置，叫什么名字都可以）中建立n个子文件夹，并把你训练数据中的n个类别的数据分别放到这n个文件夹中，每一个子文件夹都会被推断为一个新的类别。

同时class_mode这个参数要设置成"categorical", "binary", "sparse"或None之一. 该参数决定了返回的标签数组的形式, "categorical"会返回2D的one-hot编码标签,"binary"返回1D的二值标签."sparse"返回1D的整数标签,如果为None则不返回任何标签。

例如我是二分类问题，所以训练数据集文件夹中又分了两个子文件夹，利用os.walk()遍历我的原数据集，并放到这些子文件夹里,class_mode设置成'binary',生成器就会按照根据我设置的参数，循环生成1个batch的训练集，并且已经为它打好了分类标签。

接下来要做的，就是同样套路设置测试集生成器，以及将这两个生成器喂入model.fit_generator（…）训练即可