使用多个 GPU 进行高效训练的 keras#
由于实验室的电脑目前装有三张高性能的 GPU，但为了让 keras 发挥其效益，需要更改代码
当前环境为 keras 配上 tensorflow 后端，python3
三张 1080ti 显卡
我们原先的网络架构是 Sequential 架构（简单介绍）
所以每个批次都是通过 GPU 接收任务的方式执行
也就是平均分配给每个 GPU
但是显然所有的 GPU 无法充分利用
所以我们需要将批次切分成 GPU 的数量（有三片就切成三等份）
然后将所有批次的结果合并为一个输出，这样可以加快计算速度（当然每个 GPU 也会有输出限制）
使用下面的代码将 x 切分成（n_gpus) 等份，注意此时的 x 并不是真正的数字而是 tensor（即还没有运行出来的值）
def slice_batch(x, n_gpus, part):
sh = K.shape(x)
L = sh[0] // n_gpus
if part == n_gpus - 1:
return x[part * L:]
return x[part * L:(part + 1) * L]
之后将 model 传入下面的函数，就会将当前这个 model 加上（n_gpus) 个 lambda 层接在 input 上
然后将这些结果全部合并为输出（并以 batch 的 axis 合并）
def to_multi_gpu(model, n_gpus=3):
with tf.device('/cpu:0'):
print(model.input_shape)
x = Input(model.input_shape[1:], name="input")
towers = []
for g in range(n_gpus):
with tf.device('/gpu:' + str(g)):
slice_g = Lambda(slice_batch, lambda shape: shape, arguments={'n_gpus': n_gpus, 'part': g})(x)
towers.append(model(slice_g))
with tf.device('/cpu:0'):
merged = Concatenate(axis=0)(towers)
return Model(inputs=[x], outputs=[merged])
使用完的结果会类似下面的图
没有使用此技巧时训练的时间为
160 个 epoch： 23 分钟
使用此技巧平均分配到三张显卡上的时间为
160 个 epoch： 9 分钟