Tensorflow在訓練模型的時候如何指定GPU進行訓練
阿新 • • 發佈:2018-12-11
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'