使用多个 GPU 进行高效训练的 keras#
由于实验室的计算机目前安装了三个高性能的 GPU，为了让 keras 发挥其效能，需要更改代码
当前环境为 keras 配上 tensorflow 后端，python3
三张 1080ti 显卡
我们原先的网络架构是 Sequential 架构（简单介绍）
因此，每个批次都是通过 GPU 进行任务处理的方式执行
也就是将批次均匀分配给每个 GPU
但显然所有的 GPU 无法充分利用
因此，我们需要将批次切分成 GPU 的数量（有三片就切成三等份）
然后将所有批次的结果合并为一个输出，这样可以加快计算速度（当然每个 GPU 也会有输出限制）
使用下面的代码将 x 切分成（n_gpus）等份，注意此时的 x 并不是真正的数字而是张量（即尚未运行的值）
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]
之后将模型传入下面的函数，就会在当前模型上添加（n_gpus）个 lambda 层连接到输入
然后将这些结果全部合并为输出（并以批次的轴进行合并）
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 分钟