使用 predict_generator 和 VGG16 的内存错误

Question

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

#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在您的代码中）。