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

Question

我想训练一种卷积递归神经网络，用于视频帧预测。由于单个帧非常大，因此一次性将整个训练数据装入内存是具有挑战性的。因此，我遵循了一些在线教程以创建自定义数据生成器。在测试时，它似乎能正常工作，但其速度比直接使用预加载数据要慢至少100倍。由于我只能在GPU上装载大约8个批次，所以我了解到需要快速生成数据，然而这似乎并不是事实。

我在单个P100上训练我的模型，并有32GB的内存可供最多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之前准备数据似乎是瓶颈所在。

有什么想法可以提高像这样的数据生成器的性能吗？是否缺少某些东西可以保证数据及时准备好？

非常感谢！

- al_cc

---

你能解决这个问题吗？我在使用一个简单的生成器时遇到了类似的问题。获取样本需要10分钟或更长时间。 - CAta.RAy

不，我不是。但在Tensorflow 2中有一个数据类，允许构建适当的输入管道。我认为现在这将是开始的地方。 - al_cc

Answer 1

你可以尝试使用tf.data.Dataset的预取功能。预取允许你在GPU计算梯度下降的同时，使用CPU计算下一个批次或多个批次。但要注意：你需要在数据生成器中将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或更多，以获得最大速度，具体取决于您的计算机和批量大小。

- Valentin Goldité

Answer 2

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

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

- keertika jain

Answer 3

当您使用fit_generator时，可以使用workers=设置来扩展发生器工作人员的数量。但是，您应确保在getitem中考虑“item”参数，以确保不同的工作人员（未同步）根据项索引返回不同的值。即，而不是随机样本，也许只需根据索引返回数据切片。您可以在开始之前对整个数据集进行洗牌，以确保数据集顺序随机化。

- Pedro Marques

Answer 4

请尝试使用use_multiprocessing=True参数，这是我在使用您提供的数据生成器时在我的基于GTX 1080Ti的系统上观察到的数字。

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

148/148 [==============================] - 9秒 60毫秒/步

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

148/148 [==============================] - 2秒 11毫秒/步

- Manoj Mohan

---

我尝试过这个方法，但是它仍然比不使用数据生成器要慢得多。使用预加载的数据，我可以在大约10分钟内完成一个时期，而使用数据生成器（并且use_multiprocessing=True），需要大约3小时30分钟，这是无法接受的。我仍然不知道为什么会花费这么多时间，特别是因为我的生成器中没有任何花哨的东西。我已经看到了一些数据生成器，它们可以实时调整图像大小，我认为这必须比加载准备好的numpy文件要昂贵得多。 - al_cc

训练图像的数量是多少？另外，看起来您在发布代码之前进行了一些编辑，因为存在一些语法错误。也许您可以通过自己调用生成器来对代码进行分析。此外，请查看CPU利用率等信息。 - Manoj Mohan

总共有12000个培训视频，每个视频包含5帧。是的，我编辑了一下代码，以便更简洁地提出我的问题。我可以尝试使用分析工具，但仍然感到惊讶的是，为什么这个相对简单的生成器所需的时间比在线找到的自定义数据生成器执行实时图像调整还要长。 - al_cc