我们如何在 Keras 库(或 TensorFlow)中编程以在多个 GPU 上划分训练?假设您在一个具有 8 个 GPU 的 Amazon ec2 实例中,并且您希望使用所有这些 GPU 来加快训练速度,但您的代码仅适用于单个 CPU 或 GPU。
Keras 中的多 GPU
从Keras 常见问题解答中,以下是启用“数据并行”的复制粘贴代码。即让您的每个 GPU 独立处理数据的不同子集。
from keras.utils import multi_gpu_model
# Replicates `model` on 8 GPUs.
# This assumes that your machine has 8 available GPUs.
parallel_model = multi_gpu_model(model, gpus=8)
parallel_model.compile(loss='categorical_crossentropy',
optimizer='rmsprop')
# This `fit` call will be distributed on 8 GPUs.
# Since the batch size is 256, each GPU will process 32 samples.
parallel_model.fit(x, y, epochs=20, batch_size=256)
请注意,在撰写本文时,这似乎仅对 Tensorflow 后端有效。
更新(2018 年 2 月):
Keras 现在接受使用 multi_gpu_model 的自动 gpu 选择,因此您不必再对 gpu 的数量进行硬编码。此Pull Request中的详细信息。换句话说,这使代码看起来像这样:
try:
model = multi_gpu_model(model)
except:
pass
但更明确地说,您可以坚持使用以下内容:
parallel_model = multi_gpu_model(model, gpus=None)
奖金:
要检查您是否真的在使用所有 GPU,特别是 NVIDIA 的,您可以使用以下命令监控终端中的使用情况:
watch -n0.5 nvidia-smi
参考:
- 对于 TensorFlow:
这是有关如何使用的示例代码,因此为每个任务指定了带有设备/设备的列表:
# Creates a graph.
c = []
for d in ['/gpu:2', '/gpu:3']:
with tf.device(d):
a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3])
b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2])
c.append(tf.matmul(a, b))
with tf.device('/cpu:0'):
sum = tf.add_n(c)
# Creates a session with log_device_placement set to True.
sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
# Runs the op.
print(sess.run(sum))
tf 将默认使用 GPU 进行计算,即使是用于 CPU(如果存在支持的 GPU)。所以你可以只做一个for循环:“for d in ['/gpu:1','/gpu:2', '/gpu:3' ... '/gpu:8',]:” 和“tf.device(d)”应包括您的所有实例 GPU 资源。所以 tf.device() 实际上会被使用。
- 喀拉斯
对于 Keras,使用 Mxnet 而不是args.num_gpus,其中num_gpus是您所需 GPU 的列表。
def backend_agnostic_compile(model, loss, optimizer, metrics, args):
if keras.backend._backend == 'mxnet':
gpu_list = ["gpu(%d)" % i for i in range(args.num_gpus)]
model.compile(loss=loss,
optimizer=optimizer,
metrics=metrics,
context = gpu_list)
else:
if args.num_gpus > 1:
print("Warning: num_gpus > 1 but not using MxNet backend")
model.compile(loss=loss,
optimizer=optimizer,
metrics=metrics)
- horovod.tensorflow
最重要的是 Uber 最近开源了 Horovod,我认为这很棒:
import tensorflow as tf
import horovod.tensorflow as hvd
# Initialize Horovod
hvd.init()
# Pin GPU to be used to process local rank (one GPU per process)
config = tf.ConfigProto()
config.gpu_options.visible_device_list = str(hvd.local_rank())
# Build model…
loss = …
opt = tf.train.AdagradOptimizer(0.01)
# Add Horovod Distributed Optimizer
opt = hvd.DistributedOptimizer(opt)
# Add hook to broadcast variables from rank 0 to all other processes during
# initialization.
hooks = [hvd.BroadcastGlobalVariablesHook(0)]
# Make training operation
train_op = opt.minimize(loss)
# The MonitoredTrainingSession takes care of session initialization,
# restoring from a checkpoint, saving to a checkpoint, and closing when done
# or an error occurs.
with tf.train.MonitoredTrainingSession(checkpoint_dir=“/tmp/train_logs”,
config=config,
hooks=hooks) as mon_sess:
while not mon_sess.should_stop():
# Perform synchronous training.
mon_sess.run(train_op)
基本上,您可以以以下示例为例。您只需要在导入 keras 后指定 cpu 和 gpu 消耗值即可。
import keras
config = tf.ConfigProto( device_count = {'GPU': 1 , 'CPU': 56} )
sess = tf.Session(config=config)
keras.backend.set_session(sess)
之后,您将适合模型。
model.fit(x_train, y_train, epochs=epochs, validation_data=(x_test, y_test))
最后,您可以减少消耗值而不是上限工作。
我们如何使用 Horovd 和 Keras 访问多个 GPU 的简单示例:Github 代码 Keras MNIST Example with Horovod。
另外,请访问链接了解更多信息:Horovod with Keras