Keras处理超大数据集时的内存问题求解

嘿，我看到你用Keras处理1500张图时一切正常，但scale到5万张就触发了内存溢出——这个问题其实很常见，核心原因是你一次性把所有提取的特征都塞进了内存，当数据量变大时，内存自然扛不住。下面是具体的分析和解决方案：

核心问题拆解 #

你当前的代码是先用model.predict_generator把5万张图的特征全部提取出来，存在features变量里，然后再用顶层模型去预测。对于InceptionResNetV2来说，每张图的特征是(3,3,1536)的张量，5万张的话，光是这部分数据就占了大概500003315364字节≈2.7GB（按float32计算），再加上TensorFlow运行时的内存开销、生成器缓存的图像数据，很容易把主机内存（host memory）撑爆，也就是你看到的CUDA_ERROR_OUT_OF_MEMORY报错。

解决方案：分批次处理，避免一次性加载所有数据 #

我们可以把“提取特征+预测”的过程拆成小批次来做，每处理完一批就释放对应的内存，这样就能避免内存堆积。

修改后的代码 #

from __future__ import division
import numpy as np
import gc  # 用于手动回收内存
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
batch_size = 6  # 可以根据你的GPU内存适当调大，比如32或64

# 加载预训练的基础模型
base_model = applications.InceptionResNetV2(
    include_top=False, 
    weights='imagenet', 
    input_shape=(img_width, img_height, 3)
)
print("base pretrained model loaded")

# 构建图像生成器（保持shuffle=False，确保文件名和预测结果一一对应）
validation_generator = ImageDataGenerator(rescale=1./255).flow_from_directory(
    path, 
    target_size=(img_width, img_height), 
    batch_size=batch_size,
    shuffle=False
)
print("generator built")

# 加载顶层分类模型
top_model = Sequential()
top_model.add(Flatten(input_shape=(3, 3, 1536)))
top_model.add(Dense(256, activation='relu'))
top_model.add(Dropout(0.5))
top_model.add(Dense(6, activation='softmax'))
top_model.load_weights(top_model_weights_path)
print("top model loaded")

# 分批次处理：提取特征 + 预测，逐步积累结果
total_samples = validation_generator.samples
all_proba = []
all_classes = []

for step in range(total_samples // batch_size + 1):
    # 提取当前批次的特征（steps=1表示只取一个batch）
    batch_features = base_model.predict_generator(
        validation_generator, 
        steps=1, 
        verbose=1,
        workers=4,  # 用多进程加载图像，提升效率
        use_multiprocessing=True
    )
    # 用顶层模型预测当前批次的结果
    batch_proba = top_model.predict_proba(batch_features, verbose=0)
    batch_classes = top_model.predict_classes(batch_features, verbose=0)
    
    # 将批次结果加入总列表
    all_proba.extend(batch_proba)
    all_classes.extend(batch_classes)
    
    # 手动回收无用内存，避免堆积
    del batch_features, batch_proba, batch_classes
    gc.collect()

# 转换为numpy数组方便后续处理
all_proba = np.array(all_proba)
all_classes = np.array(all_classes)

print("Prediction probabilities:")
print(all_proba)
print("Predicted classes:")
print(all_classes)
print("Original file names:")
print(validation_generator.filenames)

额外的内存优化技巧 #

• 调大batch_size：如果你的GPU内存足够，把batch_size从6调到32或64，能减少批次数量，降低内存碎片化问题，同时提升处理速度
• 关闭生成器的缓存：如果不需要，可以在flow_from_directory里设置cache=False，避免生成器缓存过多图像数据
• 清理TensorFlow会话：如果运行多次实验，每次结束后可以用K.clear_session()（需要from keras import backend as K）清理TensorFlow的会话，释放GPU内存

这样修改后，你就能处理5万张甚至更多的图像，而不会出现内存溢出的问题了。

内容的提问来源于stack exchange，提问作者Rehan Aziz