如何编写高效的自定义 Keras 数据生成器

Question

我想训练一个用于视频帧预测的卷积循环神经网络。 单个帧非常大，因此一次将整个训练数据放入内存中具有挑战性。 因此，我在网上学习了一些教程来创建自定义数据生成器。 测试时，它似乎可以工作，但与直接使用预加载数据相比，它至少慢了 100 倍。 由于我只能在 GPU 上容纳大约 8 的批量大小，我知道需要非常快速地生成数据，但是，情况似乎并非如此。

我在单个 P100 上训练我的模型，并且有 32 GB 的内存可供多达 16 个内核使用。

class DataGenerator(tf.keras.utils.Sequence):

def __init__(self, images, input_images=5, predict_images=5, batch_size=16, image_size=(200, 200),
             channels=1):

    self.images = images
    self.input_images = input_images
    self.predict_images = predict_images
    self.batch_size = batch_size
    self.image_size = image_size
    self.channels = channels
    self.nr_images = int(len(self.images)-input_images-predict_images)

def __len__(self):

    return int(np.floor(self.nr_images) / self.batch_size)

def __getitem__(self, item):

    # Randomly select the beginning image of each batch
    batch_indices = random.sample(range(0, self.nr_images), self.batch_size)

    # Allocate the output images
    x = np.empty((self.batch_size, self.input_images,
                  *self.image_size, self.channels), dtype='uint8')
    y = np.empty((self.batch_size, self.predict_images,
                  *self.image_size, self.channels), dtype='uint8')

    # Get the list of input an prediction images
    for i in range(self.batch_size):
        list_images_input = range(batch_indices[i], batch_indices[i]+self.input_images)
        list_images_predict = range(batch_indices[i]+self.input_images,
                                         batch_indices[i]+self.input_images+self.predict_images)

        for j, ID in enumerate(list_images_input):
            x[i, ] = np.load(np.reshape(self.images[ID], (*self.imagesize, self.channels))

        # Read in the prediction images
        for j, ID in enumerate(list_images_predict):
            y[i, ] = np.load(np.reshape(self.images[ID], (*self.imagesize, self.channels))

    return x, y


# Training the model using fit_generator

params = {'batch_size': 8,
      'input_images': 5,
      'predict_images': 5,
      'image_size': (100, 100),
      'channels': 1
      }

data_path = "input_frames/"
input_images = sorted(glob.glob(data_path + "*.png"))
training_generator = DataGenerator(input_images, **params)

model.fit_generator(generator=training_generator, epochs=10, workers=6)

我原以为 Keras 将在 GPU 上处理当前批次时准备下一个数据批次，但它似乎没有赶上。 换句话说，在将数据发送到 GPU 之前准备数据似乎是瓶颈。

关于如何提高这样的数据生成器的性能有什么想法吗？ 是否缺少一些东西来保证及时准备数据？

非常感谢！

Answer 1

当您使用fit_generator时，有一个workers =设置可用于扩展生成器工作器的数量。 但是，您应该确保考虑getitem中的'item'参数，以确保不同的worker（未同步）根据项索引返回不同的值。 即，而不是随机样本，也许只是根据索引返回一片数据。 您可以在开始之前对整个数据集进行随机播放，以确保数据集顺序是随机的。

Answer 2

你可以试试use_multiprocessing = True吗？ 这些是我在基于GTX 1080Ti的系统上观察到的数字以及您提供的数据生成器。

model.fit_generator(generator=training_generator, epochs=10, workers=6)

148/148 [==============================] - 9s 60ms /步

model.fit_generator(generator=training_generator, epochs=10, workers=6, use_multiprocessing=True)

148/148 [==============================] - 2s 11ms /步

Answer 3

您可以尝试预取tf.data.Dataset 。 预取允许您使用 CPU 计算下一批，同时您的 GPU 计算梯度下降。 注意：你需要在数据生成器中将 numpy 数组改成tf.constant 。 然后尝试：

import tensoflow as tf

generator = DataGenerator(images)
spec = [tf.TypeSpec(shape=(generator.batch_size, generator.input_images,
                  *generator.image_size, generator.channels), dtype='uint8'),
        tf.TypeSpec(shape=(generator.batch_size, generator.predict_images,
                  *generator.image_size, generator.channels), dtype='uint8')
dataset = tf.data.Dataset.from_generator(DataGenerator, output_signature=spec)
dataset.batch(batch_size).prefetch(-1) # this order is important

# a custom training loop is better than model.fit() otherwise prefetching can fail
def train_loop(): 
    ...

您可以将prefetch()的“-1”更改为其他值，例如 1、2 或更多，以获得最大速度，具体取决于您的机器和批量大小。

Answer 4

该博客有助于使用 tf.data 设置输入数据管道，并且它比使用ImageDataGenerators更有效，并且还通过使用自定义数据目录来解释代码。 它还通过prefetch 、 cache提高了性能。

** Prefetch** 在使用当前批次时处理下一个批次。