1. 程式人生 > >keras 多GPU訓練,單GPU預測

keras 多GPU訓練,單GPU預測

多GPU訓練

keras自帶模組 multi_gpu_model,此方式為資料並行的方式,將將目標模型在多個裝置上各複製一份,並使用每個裝置上的複製品處理整個資料集的不同部分資料,最高支援在8片GPU上並行。
使用方式:

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=4) ##
parallel_model.
compile(loss='categorical_crossentropy', optimizer='adam') # 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=40, batch_size=128)

單GPU預測

因為定義的多核訓練,所以網路的每一層都是按GPU來命名的,訓練時採用多個GPU那麼當匯入引數的時候必須指定相同數量的GPU才行,如上程式碼的指定方式。但是,但我們將model切換到單GPU的環境中時,則會出現錯誤,此時我們必須將引數儲存為單GPU的形式。
方法:
在原多GPU環境中匯入模型,儲存為單GPU版本,程式碼如下:

from keras.utils import multi_gpu_model
paralleled_model=multi_gpu_model(seg_model,gpus=4)
paralleled_model.load_weights("srf_best_model.h5") # 此時model也自動載入了權重,可用model進行預測
seg_model.save('srf_single_gpu_best_model.h5')

注意:paralleled_model為多GPUmodel,而seg_model為單GPUmodel。
載入單GPU模型:

seg_model.load_weights(
"srf_single_gpu_best_model.h5")

參考: