Keras 指定多个GPU
import keras.backend.tensorflow_backend as KTF
import tensorflow as tf
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,4,5,8,9"
使用多个GPU训练
数据并行 (multi_gpu_model https://keras-zh.readthedocs.io/utils/ )
数据并行包括在每个设备上复制一次目标模型,并使用每个模型副本处理不同部分的输入数据。
Keras 有一个内置的实用函数 keras.utils.multi_gpu_model,它可以生成任何模型的数据并行版本,在多达 8 个 GPU 上实现准线性加速。
有关更多信息,请参阅 multi_gpu_model 的文档。这里是一个快速的例子:
from keras.utils import multi_gpu_model
model = Sequential()
model.add(...)
......
......
# 将 `model` 复制到 8 个 GPU 上。
# 假定你的机器有 8 个可用的 GPU。
parallel_model = multi_gpu_model(model, gpus=8)
parallel_model.compile(loss='categorical_crossentro