我想训练一种卷积递归神经网络,用于视频帧预测。由于单个帧非常大,因此一次性将整个训练数据装入内存是具有挑战性的。因此,我遵循了一些在线教程以创建自定义数据生成器。在测试时,它似乎能正常工作,但其速度比直接使用预加载数据要慢至少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之前准备数据似乎是瓶颈所在。
有什么想法可以提高像这样的数据生成器的性能吗?是否缺少某些东西可以保证数据及时准备好?
非常感谢!