使用 predict_generator 和 VGG16 的内存错误

Posted 2023-02-23

【中文标题】使用 predict_generator 和 VGG16 的内存错误

【英文标题】：Memory Error using predict_generator and VGG16

【发布时间】：2018-12-31 00:00:27

【问题描述】：

我正在尝试在我自己的数据集上应用迁移学习，该数据集存在于 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 已经以批处理形式接收数据？这里是否有人可以验证这不适用于我的系统，或者是否有其他解决方案？

【参考方案1】：

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

【讨论】：

感谢您的回复！我已将批量大小降低到 2，现在可以让它适用于 15.000 个图像（30 个类），但 30.000 似乎仍然太多。知道这是否太多了吗？

在这种情况下，当您尝试在内存中保存所有预测然后再存储它们时，您可能会用完 RAM。您可以尝试从生成器中生成，改用 predict_on_batch 并分批存储文件而不是一个大的 np 文件。