如何编写高效的自定义 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=设置可用于扩大生成器worker的数量。但是，您应该确保考虑getitem中的“item”参数，以确保不同的工作线程（未同步）根据项目索引返回不同的值。即不是随机样本，也许只是根据索引返回一部分数据。您可以在开始之前打乱整个数据集，以确保数据集顺序是随机的。