Tensorflow指定GPU進行訓練模型

實驗室共用一個深度學習伺服器，兩塊GPU,在用tensorflow訓練深度學習模型的時候，假設我們在訓練之前沒有指定GPU來進行訓練，則預設的是選用第0塊GPU來訓練我們的模型，而且其它幾塊GPU的也會顯示被佔用。有些時候，我們更希望可以通過自己指定一塊或者幾塊GPU來訓練我們的模型，而不是用這種預設的方法。接下來將簡單介紹兩種簡單的方法。

• 我們現有的GPU個數如下所示：

1.第一種是通過tf.device()函式來指定訓練時所要使用的GPU.

假設我們要用我們的第2塊GPU來訓練模型，此時可以通過下面的程式碼來指定：

tf.device('/gpu:2')
 

實驗效果如下：

我們可以看到，雖然指定了第2塊GPU來訓練，但是第一個GPU也還是被佔用，只是模型訓練的時候，是在第2塊GPU上進行。

2.通過CUDA_VISIBLE_DEVICES來指定.

同樣使用第2塊GPU來訓練模型，我們可以在我們的python程式碼中加入：

    import os
    os.environ['CUDA_VISIBLE_DEVICES']='1'

實驗效果如下：

我們可以看到，在訓練模型的時候，使用了第2塊GPU，並且其它幾塊GPU也沒有被佔，這種就相當於在我們執行程式的時候，將除第2塊以外的GPU全部遮蔽了，只有第2塊GPU對當前執行的程式是可見的。同樣，如果有不止兩塊GPU,要指定3，4塊GPU來訓練，則上面的程式碼可以改成：

os.environ['CUDA_VISIBLE_DEVICES']='3，4'

參考連結