複数の GPU を使用して効率的なトレーニングを行うための keras#
現在、研究室のコンピュータには高性能な GPU が 3 枚搭載されていますが、keras を最大限に活用するためにはコードを変更する必要があります。
現在の環境は、keras を tensorflow バックエンドと python3 で使用しています。
GPU は 3 枚の 1080ti です。
以前のネットワークアーキテクチャは Sequential アーキテクチャでした（簡単に紹介します）。
そのため、各バッチは GPU にタスクを割り当てて実行されます。
つまり、各 GPU に均等に分配されます。
しかし、全ての GPU を十分に活用することはできません。
そのため、バッチを GPU の数（3 つの場合は 3 等分）に分割し、すべてのバッチの結果を結合して 1 つの出力にすることで、計算速度を向上させることができます（もちろん、各 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）個のラムダレイヤーが追加され、入力に接続されます。
そして、これらの結果をすべて結合して出力にします（バッチの軸で結合します）。
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 エポック：23 分
このテクニックを使用して 3 つの GPU に均等に分配した場合のトレーニング時間は
160 エポック：9 分