keras 利用多張 gpu 進行有效率的訓練#
由於實驗室電腦目前裝上了三張高效能的 gpu，但為了讓 keras 跑起來符合他的效益，所以需要更改 code
目前環境為 keras 配上 tensorflow backend, python3
顯卡三張 1080ti
原先我們的網路架構為 Sequential 架構 (粗略的介紹)
所以每次的 batch 都是利用 gpu 接 task 的方式執行
也就是平均分配給每個 gpu
那想當然全部的 gpu 就沒辦法好好的善用
所以我們要把 batch slice 成 gpu 的個數（有三片就切成三等份）
再把所有 batch 結果合併起來成為一個輸出，如此就會加快運算速度（當然也會有每個 gpu 的輸出限制）
利用下面的 code 把 x 切成（n_gpus) 等份，注意此時的 x 並不是真正的數字而是 tensor（也就是還沒 run 出來的值）
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 傳入下面的 function 就會把目前這個 model 加上（n_gpus) 個 lambda layer 接上 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 分鐘