Keras数据生成器以及如何使用它们

重点介绍如何构建数据生成器以在Keras中加载和处理图像。

数据生成器的功能是什么

在Keras Model类中，有三种我们感兴趣的方法：fit_generator，valuate_generator和predict_generator。它们全部三个都需要数据生成器，但并非所有生成器都是平等创建的。
让我们看一下每种方法需要哪种生成器：

fit_generator

需要两个生成器，一个用于训练数据，另一个用于验证。幸运的是，它们两个都应该返回一个tupple（输入，目标），并且它们都可以是Sequence类的实例。

Evaluation_generator

此处的数据生成器具有与fit_generator中相同的要求，并且可以与训练生成器相同。

predict_generator

这里的生成器有点不同。它应该只返回输入。
考虑到这一点，让我们构建一些数据生成器。由于fit_generator中的生成器和evaluate_generator之间的相似性，我们将集中精力构建fit_generator和predict_generator的数据生成器。

ImageDataGenerator类

ImageDataGenerator类在图像分类中非常有用。有多种使用此生成器的方法，具体取决于我们使用的方法，这里我们将重点介绍flow_from_directory采取的路径，该目录包含在子目录中排序的图像和图像增强参数。
让我们看一个例子：
我们将使用可从https://www.kaggle.com/c/dogs-vs-cats/data下载的数据集，其结构如下：

data/
    train/
        dogs/
            dog001.jpg
            dog002.jpg
            ...
        cats/
            cat001.jpg
            cat002.jpg
            ...
    validation/
        dogs/
            dog001.jpg
            dog002.jpg
            ...
        cats/
            cat001.jpg
            cat002.jpg
            ...

首先，让我们导入所有必要的库，并创建具有图像增强功能的数据生成器。

from tensorflow.keras.preprocessing.image import ImageDataGenerator

train_datagen = ImageDataGenerator(
    rescale=1./255,
    shear_range=0.2,
    zoom_range=0.2,
    horizontal_flip=True)

test_datagen = ImageDataGenerator(rescale=1./255)

train_generator = train_datagen.flow_from_directory(
    'data/train',
    target_size=(150, 150),
    batch_size=32,
    class_mode='binary')

validation_generator = test_datagen.flow_from_directory(
    'data/validation',
    target_size=(150, 150),
    batch_size=32,
    class_mode='binary')

最后，创建一个模型并运行fit_generator方法。

model.fit_generator(
    train_generator,
    steps_per_epoch=2000,
    epochs=50,
    validation_data=validation_generator,
    validation_steps=800)

灵活的数据生成器

要构建自定义数据生成器，我们需要从 Sequence 类继承。让我们添加所需的参数。

class DataGenerator(Sequence):
    'Generates data for Keras'
    def __init__(self, list_IDs, labels, image_path, mask_path,
    to_fit=True, batch_size=32, dim=(256,256),
    n_channels=1, n_classes=10, shuffle=True):
        'Initialization'
        self.list_IDs = list_IDs
        self.labels = labels
        self.image_path = image_path
        self.mask_path = mask_path
        self.to_fit = to_fit
        self.batch_size = batch_size
        self.dim = dim
        self.n_channels = n_channels
        self.n_classes = n_classes
        self.shuffle = shuffle
        self.on_epoch_end()

Sequence类迫使我们实现两种方法：len__和__getitem。如果我们希望生成器在每个时期之后执行某些操作，我们还可以实现on_epoch_end方法。
__len__方法应返回每个时期的批次数。下面显示了一种可能的实现。

def __len__(self):
    'Denotes the number of batches per epoch'
    return int(np.floor(len(self.list_IDs) / self.batch_size))

如果shuffle = True，此示例中的on_epoch_end可以将训练的索引洗牌。但是，在每个epoch之后，我们可以运行任何逻辑。

def on_epoch_end(self):
    'Updates indexes after each epoch'
    self.indexes = np.arange(len(self.list_IDs))
    if self.shuffle == True:
        np.random.shuffle(self.indexes)

我们必须实现的第二种方法是__getitem__，它完全符合您的期望。如果我们预测的话，它应该返回一批图像和蒙版。可以通过将to_fit设置为True或False来控制。

def __getitem__(self, index):
    'Generate one batch of data'
    # Generate indexes of the batch
    indexes = self.indexes[index*self.batch_size: (index+1)*self.batch_size]
    
    # Find list of IDs
    list_IDs_temp = [self.list_IDs[k] for k in indexes]
    
    # Generate data
    X = self._generate_X(list_IDs_temp)
  
    if self.to_fit:
        y = self._generate_y(list_IDs_temp)
        return X, y
    else
        return X

def _generate_X(self, list_IDs_temp):
    'Generates data containing batch_size images'
    # Initialization
    X = np.empty((self.batch_size, *self.dim, self.n_channels))
    
    # Generate data
    for i, ID in enumerate(list_IDs_temp):
        # Store sample
        X[i,] = _load_grayscale_image(self.image_path + self.labels[ID])
        
    return X
 
def _generate_y(self, list_IDs_temp):
    'Generates data containing batch_size masks'
    y = np.empty((self.batch_size, *self.dim), dtype=int)
    
    # Generate data
    for i, ID in enumerate(list_IDs_temp):
        # Store sample
        y[i,] = _load_grayscale_image(self.mask_path + self.labels[ID])
        
    return y
 
def _load_grayscale_image(image_path):
    'Load grayscale image'
    img = cv2.imread(image_path)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    img = img / 255
    
    return img

整个数据生成器应与此类似：

import numpy as np
import cv2
from tensorflow.keras.utils import Sequence


class DataGenerator(Sequence):
    """Generates data for Keras
    Sequence based data generator. Suitable for building data generator for training and prediction.
    """
    def __init__(self, list_IDs, labels, image_path, mask_path,
                 to_fit=True, batch_size=32, dim=(256, 256),
                 n_channels=1, n_classes=10, shuffle=True):
        """Initialization
        :param list_IDs: list of all 'label' ids to use in the generator
        :param labels: list of image labels (file names)
        :param image_path: path to images location
        :param mask_path: path to masks location
        :param to_fit: True to return X and y, False to return X only
        :param batch_size: batch size at each iteration
        :param dim: tuple indicating image dimension
        :param n_channels: number of image channels
        :param n_classes: number of output masks
        :param shuffle: True to shuffle label indexes after every epoch
        """
        self.list_IDs = list_IDs
        self.labels = labels
        self.image_path = image_path
        self.mask_path = mask_path
        self.to_fit = to_fit
        self.batch_size = batch_size
        self.dim = dim
        self.n_channels = n_channels
        self.n_classes = n_classes
        self.shuffle = shuffle
        self.on_epoch_end()

    def __len__(self):
        """Denotes the number of batches per epoch
        :return: number of batches per epoch
        """
        return int(np.floor(len(self.list_IDs) / self.batch_size))

    def __getitem__(self, index):
        """Generate one batch of data
        :param index: index of the batch
        :return: X and y when fitting. X only when predicting
        """
        # Generate indexes of the batch
        indexes = self.indexes[index * self.batch_size:(index + 1) * self.batch_size]

        # Find list of IDs
        list_IDs_temp = [self.list_IDs[k] for k in indexes]

        # Generate data
        X = self._generate_X(list_IDs_temp)

        if self.to_fit:
            y = self._generate_y(list_IDs_temp)
            return X, y
        else:
            return X

    def on_epoch_end(self):
        """Updates indexes after each epoch
        """
        self.indexes = np.arange(len(self.list_IDs))
        if self.shuffle == True:
            np.random.shuffle(self.indexes)

    def _generate_X(self, list_IDs_temp):
        """Generates data containing batch_size images
        :param list_IDs_temp: list of label ids to load
        :return: batch of images
        """
        # Initialization
        X = np.empty((self.batch_size, *self.dim, self.n_channels))

        # Generate data
        for i, ID in enumerate(list_IDs_temp):
            # Store sample
            X[i,] = self._load_grayscale_image(self.image_path + self.labels[ID])

        return X

    def _generate_y(self, list_IDs_temp):
        """Generates data containing batch_size masks
        :param list_IDs_temp: list of label ids to load
        :return: batch if masks
        """
        y = np.empty((self.batch_size, *self.dim), dtype=int)

        # Generate data
        for i, ID in enumerate(list_IDs_temp):
            # Store sample
            y[i,] = self._load_grayscale_image(self.mask_path + self.labels[ID])

        return y

    def _load_grayscale_image(self, image_path):
        """Load grayscale image
        :param image_path: path to image to load
        :return: loaded image
        """
        img = cv2.imread(image_path)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        img = img / 255
        return img

假设我们有两个目录，一个目录保存图像，另一个目录保存遮罩图像，并且每个图像都有一个具有相同名称的对应遮罩，以下代码将使用自定义数据生成器训练模型。

image_path = 'path to images'
mask_path = 'path to masks'

training_generator = DataGenerator(train_idx, labels, image_path, mask_path)
validation_generator = DataGenerator(val_idx, labels, image_path, mask_path)

# Design model
model = Sequential()
[...] # Architecture
model.compile()

# Train model on dataset
model.fit(training_generator, validation_data=validation_generator)

最后，如果我们要使用数据生成器进行预测，则应将to_fit设置为False并应调用predict_generator。

image_path = 'path to images'

pred_labels = [...] # list of image names

pred_generator = DataGenerator(pred_idx, pred_labels, image_path, to_fit=False)

pred = model.predict_generator(pred_generator)

可运行实例

import numpy as np
import keras
from keras.layers import Dense


class DataGenerator(keras.utils.Sequence):
    'Generates data for Keras'

    def __init__(self, train_data, test_data, batch_size=32, n_channels=1, shuffle=True):
        'Initialization'
        self.train_data = train_data
        self.test_data = test_data
        self.batch_size = batch_size
        self.n_channels = n_channels
        self.shuffle = shuffle
        self.on_epoch_end()

    def __len__(self):
        'Denotes the number of batches per epoch'
        return int(np.floor(len(self.train_data) / self.batch_size))

    def __getitem__(self, index):
        'Generate one batch of data'
        # Generate indexes of the batch
        indexes = self.indexes[index * self.batch_size:(index + 1) * self.batch_size]

        # Generate data
        # X, y = self.__data_generation(train_data_temp)

        X = self.train_data[indexes]
        y = self.test_data[indexes]

        return X, y

    def on_epoch_end(self):
        'Updates indexes after each epoch'
        self.indexes = np.arange(len(self.train_data))
        if self.shuffle == True:
            np.random.shuffle(self.indexes)


def gen_model():

    model = keras.models.Sequential()

    model.add(Dense(units=1, use_bias=False, input_shape=(1,)))  # 仅有的1个权重在这里

    return model


if __name__ == '__main__':

    # 数据比较简单，用 CPU 即可
    import os
    os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

    # Parameters
    params = {'batch_size': 10,
              'n_channels': 1,
              'shuffle': False}

    # Datasets
    # partition =  # IDs
    # labels =  # Labels
    x = {}
    x['train'] = np.arange(100, dtype='int32')
    x['validation'] = np.arange(100, 120, dtype='int32')
    y = {}
    y['train'] = -x['train']
    y['test'] = -x['train']


    # Generators
    training_generator = DataGenerator(x['train'], y['train'], **params)
    # validation_generator = DataGenerator(x['validation'], y['test'], **params)

    # Design model
    model = gen_model()
    model.compile(loss='mse', optimizer='adam')

    # model.fit(x['train'], y['train'], epochs=1000, batch_size=10, verbose=2)

    # Train model on dataset
    model.fit_generator(generator=training_generator,
                        # validation_data=validation_generator,
                        epochs=2000,
                        verbose=2,
                        workers=6,
                        )

    print(model.layers[0].get_weights())

结论

尽管Keras提供了数据生成器，但它们的功能有限。原因之一是每个任务都需要一个不同的数据加载器。有时每张图像都有一个遮罩，有时是几个，有时将遮罩保存为图像，有时将其编码，等等。
对于每个任务，我们可能都需要调整数据生成器，但结构将保持不变。

---

ref：

文中代码的github地址
Keras data generators and how to use them

相关阅读：

阅读以下的链接，你能够更加清晰地理解。
姊妹文章：如何在Keras中使用数据生成器（data generators）的详细示例

A detailed example of how to use data generators with Keras by Stanford

Keras Model class

Keras ImageDataGenerator class