Tensorflow中怎么利用CUDA_VISIBLE_DEVICES控制GPU

本篇文章給大家分享的是有關(guān)Tensorflow中怎么利用CUDA_VISIBLE_DEVICES控制GPU，小編覺得挺實(shí)用的，因此分享給大家學(xué)習(xí)，希望大家閱讀完這篇文章后可以有所收獲，話不多說，跟著小編一起來看看吧。

• os.environ[“CUDA_DEVICE_ORDER”] = “PCI_BUS_ID” # 按照PCI_BUS_ID順序從0開始排列GPU設(shè)備

• os.environ[“CUDA_VISIBLE_DEVICES”] = “0” #設(shè)置當(dāng)前使用的GPU設(shè)備僅為0號(hào)設(shè)備 設(shè)備名稱為’/gpu:0’

• os.environ[“CUDA_VISIBLE_DEVICES”] = “1” #設(shè)置當(dāng)前使用的GPU設(shè)備僅為1號(hào)設(shè)備 設(shè)備名稱為’/gpu:1’

• os.environ[“CUDA_VISIBLE_DEVICES”] = “0,1” #設(shè)置當(dāng)前使用的GPU設(shè)備為0,1號(hào)兩個(gè)設(shè)備,名稱依次為’/gpu:0’、’/gpu:1’

• os.environ[“CUDA_VISIBLE_DEVICES”] = “1,0” #設(shè)置當(dāng)前使用的GPU設(shè)備為1,0號(hào)兩個(gè)設(shè)備,名稱依次為’/gpu:1’、’/gpu:0’。表示優(yōu)先使用1號(hào)設(shè)備,然后使用0號(hào)設(shè)備

如果服務(wù)器有多個(gè)GPU，tensorflow默認(rèn)會(huì)全部使用。如果只想使用部分GPU，可以通過參數(shù)CUDA_VISIBLE_DEVICES來設(shè)置GPU的可見性。

示例：

Environment Variable Syntax      Results
CUDA_VISIBLE_DEVICES=1           Only device 1 will be seen
CUDA_VISIBLE_DEVICES=0,1         Devices 0 and 1 will be visible
CUDA_VISIBLE_DEVICES="0,1">

在終端設(shè)置

在終端調(diào)用Python腳本前，可以設(shè)置CUDA_VISIBLE_DEVICES變量，如下：

$ CUDA_VISIBLE_DEVICES=1 python my_script.py

這樣my_script.py腳本就只能使用GPU 1。

在Python腳本內(nèi)設(shè)置

如果想在Python的腳本內(nèi)設(shè)置使用的GPU，可以使用os.environ，如下：

import os
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"   # see issue #152
os.environ["CUDA_VISIBLE_DEVICES"]="1"

檢查TensorFlow對(duì)GPU的可見性：

from tensorflow.python.client import device_lib
print(device_lib.list_local_devices())

以上就是Tensorflow中怎么利用CUDA_VISIBLE_DEVICES控制GPU，小編相信有部分知識(shí)點(diǎn)可能是我們?nèi)粘９ぷ鲿?huì)見到或用到的。希望你能通過這篇文章學(xué)到更多知識(shí)。更多詳情敬請(qǐng)關(guān)注億速云行業(yè)資訊頻道。