我用的 keras 提供的 multi_gpu_model,用法参考了 keras 的官方文档:
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)
但就我的应用(一个时间序列 LSTM 模型)而言,用两块 GPU 运算单轮训练的时间甚至还稍高于用单块 GPU。
现在 multi_gpu_model 功能还没有正式 release,网上公开信息不多。有人用过这个吗,来谈谈优化方法吧。