Tensorflow中怎么利用CUDA_VISIBLE_DEVICES控制GPU

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

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

• os.environ[“CUDA_VISIBLE_DEVICES”] = “0” #設置當前使用的GPU設備僅為0號設備 設備名稱為’/gpu:0’

• os.environ[“CUDA_VISIBLE_DEVICES”] = “1” #設置當前使用的GPU設備僅為1號設備 設備名稱為’/gpu:1’

• os.environ[“CUDA_VISIBLE_DEVICES”] = “0,1” #設置當前使用的GPU設備為0,1號兩個設備,名稱依次為’/gpu:0’、’/gpu:1’

• os.environ[“CUDA_VISIBLE_DEVICES”] = “1,0” #設置當前使用的GPU設備為1,0號兩個設備,名稱依次為’/gpu:1’、’/gpu:0’。表示優先使用1號設備,然后使用0號設備

如果服務器有多個GPU，tensorflow默認會全部使用。如果只想使用部分GPU，可以通過參數CUDA_VISIBLE_DEVICES來設置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">

在終端設置

在終端調用Python腳本前，可以設置CUDA_VISIBLE_DEVICES變量，如下：

$ CUDA_VISIBLE_DEVICES=1 python my_script.py

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

在Python腳本內設置

如果想在Python的腳本內設置使用的GPU，可以使用os.environ，如下：

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

檢查TensorFlow對GPU的可見性：

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

以上就是Tensorflow中怎么利用CUDA_VISIBLE_DEVICES控制GPU，小編相信有部分知識點可能是我們日常工作會見到或用到的。希望你能通過這篇文章學到更多知識。更多詳情敬請關注億速云行業資訊頻道。