Keras 指定多个GPU，使用多个GPU训练

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_crossentropy',
                       optimizer='rmsprop')

# 这个 `fit` 调用将分布在 8 个 GPU 上。
# 由于 batch size 为 256，每个 GPU 将处理 32 个样本。
parallel_model.fit(x, y, epochs=20, batch_size=256)

保存模型

关于模型保存

要保存多 GPU 模型，请通过模板模型（传递给 multi_gpu_model 的参数）调用 .save(fname) 或 .save_weights(fname) 以进行存储，而不是通过 multi_gpu_model 返回的模型。

model.save(...)  而不是  parallel_model.save(...)

GPU动态增长，不全部占满显存, 按需分配

import keras.backend.tensorflow_backend as KTF
import tensorflow as tf
import os
 
 
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
 
config = tf.ConfigProto()
config.gpu_options.allow_growth=True   #不全部占满显存, 按需分配
sess = tf.Session(config=config)
 
KTF.set_session(sess)

 

可能会有一个报错：TypeError: can’t pickle thread.lock objects

解决办法：自定义一个（ModelCheckpoint检查点）回调函数

Reference：https://blog.csdn.net/z5217632/article/details/80952372