Keras多GPU环境下运行predict_generator生成VGG16瓶颈特征的问题

哎，这个坑我之前也踩过！Keras 2.0.9的multi_gpu_model在处理predict_generator的时候确实有点“偷懒”——它虽然会把模型复制到各个GPU上，但默认实现会把所有预测结果先拉回CPU再合并，导致你看着GPU占用了，但实际计算全在CPU跑，速度慢得离谱。

不过不用太担心要重写整个predict_generator，有几个简便的小技巧能搞定这个问题：

1. 调整multi_gpu_model参数，关闭CPU合并 #

首先，调用multi_gpu_model时加上cpu_merge=False参数，强制让计算留在GPU上：

from keras.utils import multi_gpu_model

parallel_model = multi_gpu_model(your_vgg16_model, gpus=6, cpu_merge=False)

这个参数会让模型不再自动把GPU的计算结果拉回CPU合并，而是直接返回每个GPU的输出列表，从根源上避免CPU兜底计算。

2. 包装生成器，适配多GPU输入 #

因为cpu_merge=False时，模型期望输入是对应GPU数量的批量列表（每个元素对应一块GPU的输入数据），所以你需要给原始生成器套一层包装：

def multi_gpu_data_generator(original_generator, gpus=6):
    while True:
        # 每次获取gpus个批量的数据
        batches = [next(original_generator) for _ in range(gpus)]
        # 如果生成器只返回输入数据（比如生成瓶颈特征不需要标签），直接用batches即可
        # 如果生成器返回(inputs, labels)，就提取inputs部分
        input_batches = [batch[0] for batch in batches]
        yield input_batches

这个包装后的生成器每次会吐出6个批量，分别喂给6块GPU，完美适配多GPU模型的输入要求。

3. 运行预测并合并结果 #

接下来调用predict_generator时，记得调整steps_per_epoch——现在每个step会处理6个批量，所以步数要改成原来的1/6（如果原来的步数是N，现在就用N // 6，保证总数据量一致）：

# 假设原来的steps_per_epoch是total_samples // batch_size
new_steps = original_steps // 6
gpu_outputs = parallel_model.predict_generator(
    multi_gpu_data_generator(your_original_generator),
    steps=new_steps
)

# 合并所有GPU的输出结果，得到完整的瓶颈特征数组
final_features = np.concatenate(gpu_outputs, axis=0)

额外注意事项 #

• 确保原始生成器的批量大小设置合理，避免单GPU批量过大导致显存溢出
• 检查TensorFlow 1.x（Keras 2.0.9对应TF1.x版本）是否正确识别了所有6块GPU，可以用tf.test.is_gpu_available()确认
• 如果数据集总数不是6的倍数，最后可能会有少量剩余数据，记得单独用单GPU处理这部分

亲测这个方法在Keras 2.0.9+TF1.12环境下有效，不用重写核心预测逻辑，只是做了一点包装和参数调整，应该能解决你的问题。

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