keras利用多張gpu進行快速訓練

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分鐘

​