Keras+TensorFlow GPU处理大规模图像数据时内存溢出问题求助

看起来你遇到的核心问题是生成bottleneck特征时，内存（包括主机内存和TensorFlow的固定内存）被耗尽——尤其是你尝试一次性处理120万张图像的特征，直接把所有结果存在内存里，这显然超出了系统的承载能力。结合你的硬件配置和错误信息，我给你以下几个针对性的解决方案：

1. 分批生成并增量保存Bottleneck特征（最关键） #

你当前的代码是一次性生成所有训练集的bottleneck特征，然后整体保存到内存中的bottleneck_features_train变量——120万张图的VGG16特征（每张是7×7×512的float32张量），总大小约为120GB，这已经接近你的128GB内存上限，再加上系统和其他进程的占用，必然会触发内存错误。

改用HDF5格式增量写入，避免一次性加载所有数据到内存：

import numpy as np
from keras.preprocessing.image import ImageDataGenerator
from keras.models import Sequential
from keras.layers import Dropout, Flatten, Dense
from keras.applications import VGG16
from keras.utils import HDF5Matrix, to_categorical
import h5py
import tensorflow as tf
from keras.backend.tensorflow_backend import set_session

# 先配置TensorFlow会话，优化内存使用
config = tf.ConfigProto()
config.gpu_options.allow_growth = True  # GPU按需分配内存，避免一次性占满显存
config.gpu_options.pinned_memory = False  # 关闭固定主机内存分配，解决报错中的pinned memory问题
set_session(tf.Session(config=config))

# 图像参数
img_width, img_height = 224, 224
top_model_weights_path = 'bottleneck_fc_model.h5'
train_data_dir = 'data/train'
validation_data_dir = 'data/validation'
nb_train_samples = 1200000
nb_validation_samples = 60000
epochs = 50
batch_size = 10
num_classes = 15000  # 匹配你的数据集类别数

# 构建VGG16模型，不包含顶部全连接层
model_vgg = VGG16(include_top=False, weights='imagenet')

# 生成训练集bottleneck特征——分批写入HDF5
datagen = ImageDataGenerator(rescale=1. / 255)

train_generator_bottleneck = datagen.flow_from_directory(
    train_data_dir,
    target_size=(img_width, img_height),
    batch_size=batch_size,
    class_mode=None,
    shuffle=False,
    workers=8,  # 利用多进程加载图像，适配你的16核CPU
    use_multiprocessing=True,
    max_queue_size=10  # 限制预加载队列大小，减少内存占用
)

# 创建HDF5文件存储特征和标签
with h5py.File('bottleneck_features_train.h5', 'w') as f:
    # 创建特征数据集，预设形状和类型
    feat_dset = f.create_dataset('features', shape=(nb_train_samples, 7, 7, 512), dtype='float32')
    # 同时保存one-hot编码的标签
    label_dset = f.create_dataset('labels', shape=(nb_train_samples, num_classes), dtype='float32')
    
    total_batches = nb_train_samples // batch_size
    for batch_idx in range(total_batches):
        # 获取当前批次图像
        x_batch = next(train_generator_bottleneck)
        # 预测特征
        feat_batch = model_vgg.predict_on_batch(x_batch)
        # 获取当前批次的标签并转one-hot
        label_batch = to_categorical(
            train_generator_bottleneck.classes[batch_idx*batch_size : (batch_idx+1)*batch_size],
            num_classes=num_classes
        )
        # 写入HDF5文件
        start = batch_idx * batch_size
        end = start + batch_size
        feat_dset[start:end] = feat_batch
        label_dset[start:end] = label_batch
        
        # 打印进度，方便监控
        if batch_idx % 100 == 0:
            print(f"Processed {start}/{nb_train_samples} images")

# 验证集用同样的方式处理
validation_generator_bottleneck = datagen.flow_from_directory(
    validation_data_dir,
    target_size=(img_width, img_height),
    batch_size=batch_size,
    class_mode=None,
    shuffle=False,
    workers=8,
    use_multiprocessing=True,
    max_queue_size=10
)

with h5py.File('bottleneck_features_validation.h5', 'w') as f:
    feat_dset = f.create_dataset('features', shape=(nb_validation_samples, 7, 7, 512), dtype='float32')
    label_dset = f.create_dataset('labels', shape=(nb_validation_samples, num_classes), dtype='float32')
    
    total_batches = nb_validation_samples // batch_size
    for batch_idx in range(total_batches):
        x_batch = next(validation_generator_bottleneck)
        feat_batch = model_vgg.predict_on_batch(x_batch)
        label_batch = to_categorical(
            validation_generator_bottleneck.classes[batch_idx*batch_size : (batch_idx+1)*batch_size],
            num_classes=num_classes
        )
        start = batch_idx * batch_size
        end = start + batch_size
        feat_dset[start:end] = feat_batch
        label_dset[start:end] = label_batch

# 训练顶部全连接模型——用HDF5Matrix加载，无需一次性读入内存
train_data = HDF5Matrix('bottleneck_features_train.h5', 'features')
train_labels = HDF5Matrix('bottleneck_features_train.h5', 'labels')

validation_data = HDF5Matrix('bottleneck_features_validation.h5', 'features')
validation_labels = HDF5Matrix('bottleneck_features_validation.h5', 'labels')

model_top = Sequential()
model_top.add(Flatten(input_shape=(7,7,512)))
model_top.add(Dense(256, activation='relu'))
model_top.add(Dropout(0.5))
model_top.add(Dense(num_classes, activation='softmax'))  # 多分类场景用softmax

model_top.compile(optimizer='rmsprop',
                  loss='categorical_crossentropy',  # 匹配多分类的损失函数
                  metrics=['accuracy'])

model_top.fit(train_data, train_labels,
              epochs=epochs,
              batch_size=batch_size,
              validation_data=(validation_data, validation_labels))

model_top.save_weights(top_model_weights_path)

2. 关闭TensorFlow固定主机内存分配 #

你的错误日志里明确提到could not allocate pinned host memory，TensorFlow默认会使用固定主机内存来加速GPU-CPU数据传输，但当内存不足时可以关闭这个选项——上面的代码已经包含了这个配置，核心逻辑是：

config = tf.ConfigProto()
config.gpu_options.pinned_memory = False
set_session(tf.Session(config=config))

3. 修正标签生成逻辑 #

你原来的代码是按二分类生成标签（[0]*(nb//2)+[1]*(nb//2)），但你的数据集有15000个类别，这明显不匹配。改用flow_from_directory的classes属性获取真实的类别索引，再转成one-hot编码，上面的代码已经做了修正。

4. 优化系统交换空间利用率 #

虽然你有300GB交换分区，但系统默认可能不会优先使用它。可以临时调整swappiness参数：

sudo sysctl vm.swappiness=100

如果要永久生效，编辑/etc/sysctl.conf，添加一行vm.swappiness=100，然后重启系统。不过交换空间速度远慢于内存，这是最后的兜底方案。

5. 确认CUDA与TensorFlow版本兼容性 #

CUDA 9.1对应的TensorFlow GPU版本应该是1.8.x或1.9.x，如果你的TensorFlow版本过高或过低，可能会出现内存分配异常。可以用以下命令检查版本：

pip show tensorflow-gpu

如果版本不匹配，升级或降级到兼容版本即可。

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