Keras多GPU及分布式

如何在多张GPU卡上使用Keras？

我们建议有多张GPU卡可用时，使用TnesorFlow后端。
有两种方法可以在多张GPU上运行一个模型：数据并行/设备并行
大多数情况下，你需要的很可能是“数据并行”

数据并行

数据并行将目标模型在多个设备上各复制一份，并使用每个设备上的复制品处理整个数据集的不同部分数据。Keras在 keras.utils.multi_gpu_model 中提供有内置函数，该函数可以产生任意模型的数据并行版本，最高支持在8片GPU上并行。请参考utils中的multi_gpu_model文档。下面是一个例子：

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)

数据并行利用多块GPU同时训练多个batch数据，运行在每块GPU上的模型为同一个神经网络，网络结构完全一样，并且共享模型参数。

设备并行

设备并行是在不同设备上运行同一个模型的不同部分，当模型含有多个并行结构，例如含有两个分支时，这种方式很适合。
这种并行方法可以通过使用TensorFlow device scopes实现，下面是一个例子：

# Model where a shared LSTM is used to encode two different sequences in parallel
input_a = keras.Input(shape=(140, 256))
input_b = keras.Input(shape=(140, 256))

shared_lstm = keras.layers.LSTM(64)

# Process the first sequence on one GPU
with tf.device_scope('/gpu:0'):
    encoded_a = shared_lstm(tweet_a)
# Process the next sequence on another GPU
with tf.device_scope('/gpu:1'):
    encoded_b = shared_lstm(tweet_b)

# Concatenate results on CPU
with tf.device_scope('/cpu:0'):
    merged_vector = keras.layers.concatenate([encoded_a, encoded_b],
                                             axis=-1)

多任务输出数据并行

在Keras版的Faster-RCNN，它由多个输出支路，也就是多个loss，在网络定义的时候一般会给命名，然后编译的时候找到不同支路layer的名字即可，就像这样：

model.compile(optimizer=optimizer, 
              loss={'main_output': jaccard_distance_loss, 'aux_output': 'binary_crossentropy'},
              metrics={'main_output': jaccard_distance_loss, 'aux_output': 'acc'},
              loss_weights={'main_output': 1., 'aux_output': 0.5})

其中main_output和aux_output就是认为定义的layer name，但是如果用了keras.utils.training_utils.multi_gpu_model()以后，名字就自动换掉了，变成默认的concatenate_1, concatenate_2等等，因此你需要先model.summary()一下，打印出来网络结构，然后弄明白哪个输出代表哪个支路，然后重新编译网络，如下：

from keras.optimizers import Adam, RMSprop, SGD
model.compile(optimizer=RMSprop(lr=0.045, rho=0.9, epsilon=1.0), 
              loss={'concatenate_1': jaccard_distance_loss, 'concatenate_2': 'binary_crossentropy'},
              metrics={'concatenate_1': jaccard_distance_loss, 'concatenate_2': 'acc'},
              loss_weights={'concatenate_1': 1., 'concatenate_2': 0.5})

而且在在Keras版的Faster-RCNN中，每个batch里，对RPN进行训练，测试后的结果作为检测网络的输入，来训练，最后把2个模型对参数的训练结果作为一个模型保存下来。

分布式

keras的分布式是利用TensorFlow实现的，要想完成分布式的训练，你需要将Keras注册在连接一个集群的TensorFlow会话上：

server = tf.train.Server.create_local_server()
sess = tf.Session(server.target)

from keras import backend as K
K.set_session(sess)