python - 使用predict_generator和VGG16的内存错误

Question

我正在尝试在我自己的数据集上应用迁移学习，该数据集包含 33.000 个训练图像(总共 1.4GB)。在使用 Keras (2.2.0) 中的 Predict_generator 进行预测时，我遇到了内存错误。当查看我的任务管理器时，我可以看到内存慢慢达到 Tesla K80 (1GPU) 的最大 VRAM 5GB。我正在使用以下代码:

#Train
print('train dataset:')
datagen = ImageDataGenerator(preprocessing_function=preprocess_input)
train_generator = datagen.flow_from_directory(
    train_data_dir,
    target_size=(img_width, img_height),
    batch_size=batch_size,
    class_mode=None,
    shuffle=False)

num_classes = len(train_generator.class_indices)
nb_train_samples = len(train_generator.filenames)
predict_size_train = int(math.ceil(nb_train_samples / batch_size))
VGG16_bottleneck_features_train = model.predict_generator(train_generator, predict_size_train, verbose=1)
np.save('XVGG16_bottleneck_features_train.npy', VGG16_bottleneck_features_train)

我尝试了很多东西，但似乎无法让它为我工作。我读过许多建议使用批处理的解决方案，但我认为我的 Predict_generator 已经以批处理形式接收数据？这里有人可以验证这对我的系统不起作用还是有其他可能的解决方案？

Answer 1

推理所需的内存与您尝试同时通过网络传递的图像数量(即批量大小)成正比。您可以尝试较小的批量大小，直到它运行为止。批量大小越小，生成器需要生成的批量越多才能传递整个数据集(代码中的 predict_size_train)。