我有这个使用 Keras 2.0.1 的代码，它是通过执行深度学习来进行图像识别的生成器。这段代码目前在 GPU 上运行时可以处理 1500 张图像，但是当我开始使用 50k 图像评估它时，我会遇到内存问题。我使用目录中的流来读取图像，使用 predict_generator 来获取预测和概率。以下是我得到的错误和我正在使用的代码：
错误：
2018-01-08 12:28:49.940361: E tensorflow/stream_executor/cuda/cuda_driver.cc:955] failed to alloc 17179869184 bytes on host: CUDA_ERROR_OUT_OF_MEMORY
2018-01-08 12:28:49.968880: W ./tensorflow/core/common_runtime/gpu/pool_allocator.h:195] could not allocate pinned host memory of size: 17179869184
代码
from __future__ import division
import numpy as np
from keras import applications
from keras.preprocessing.image import ImageDataGenerator
from keras.models import Sequential
from keras.layers import Dense, Dropout, Flatten
top_model_weights_path = '/home/rehan/ethnicity.071217.23-0.28.hdf5'
path = "/home/rehan/countries/pakistan/guys/test/"
img_width, img_height = 139, 139
confidence = 0.8
model = applications.InceptionResNetV2(include_top=False, weights='imagenet',
input_shape=(img_width, img_height, 3))
print("base pretrained model loaded")
validation_generator = ImageDataGenerator(rescale=1./255).flow_from_directory(path, target_size=(img_width, img_height),
batch_size=6,shuffle=False)
print("generator built")
features = model.predict_generator(validation_generator)
print("features found")
model = Sequential()
model.add(Flatten(input_shape=(3, 3, 1536)))
model.add(Dense(256, activation='relu'))
model.add(Dropout(0.5))
model.add(Dense(6, activation='softmax'))
model.load_weights(top_model_weights_path)
print("top model loaded")
prediction_proba = model.predict_proba(features)
prediction_classes = model.predict_classes(features)
print(prediction_proba)
print(prediction_classes)
print("original file names")
print(validation_generator.filenames)