1. 程式人生 > >python&Keras實現多GPU或指定GPU的使用

python&Keras實現多GPU或指定GPU的使用

1. keras新版本中加入多GPU並行使用的函式
下面程式段即可實現一個或多個GPU加速: 
注意:使用多GPU加速時,Keras版本必須是Keras2.0.9以上版本

from keras.utils.training_utils import multi_gpu_model   #匯入keras多GPU函式
import VGG19     #匯入已經寫好的函式模型,例如VGG19

if G <= 1:
    print("[INFO] training with 1 GPU...")
    model = VGG19()

# otherwise, we are compiling using multiple GPUs
else:
    print("[INFO] training with {} GPUs...".format(G))
    # we'll store a copy of the model on *every* GPU and then combine
    # the results from the gradient updates on the CPU
    with tf.device("/cpu:0"):
        # initialize the model
        model1 = VGG19()
        # make the model parallel(if you have more than 2 GPU)
        model = multi_gpu_model(model1, gpus=G)

2.指定使用某個GPU
首先在終端檢視主機中GPU編號:

watch -n -9 nvidia-smi

查詢結果如下所示: 
 
顯示主機中只有一塊GPU,編號為0

2.1 下面方法是直接在終端執行時加入相關語句實現指定GPU的使用
export CUDA_VISIBLE_DEVICES=0 python test.py
# 表示執行test.py檔案時,使用編號為0的GPU卡
export CUDA_VISIBLE_DEVICES=0,2 python test.py
# 表示執行test.py檔案時,使用編號為0和2的GPU卡

2.2 下面方法是在Python程式中新增
import os
# 使用第一張與第三張GPU卡
os.environ["CUDA_VISIBLE_DEVICES"] = "0, 2"