筆者需要tensorflow僅運行在一個GPU上（機器本身有多GPU），而且需要依據系統參數動態調節，故無法簡單使用CUDA_VISIBLE_DEVICES。
一種方式是全局使用tf.device函數生成的域，但設備號需要在繪制Graph前指定，仍然不夠靈活。
查閱文檔發現config的GPUOptions中的visible_device_list可以定義GPU編號從visible到virtual的映射，即可以設置tensorflow可見的GPU device，從而全局設置了tensorflow可見的GPU編號。代碼如下：
|
1
2
3
|
config = tf.ConfigProto()config.gpu_options.visible_device_list = str(device_num)sess = tf.Session(config=config) |
參考 多卡服務器下隱藏部分 GPU 和 TensorFlow 的顯存使用設置，還可以通過os包設置全局變量CUDA_VISIBLE_DEVICES，代碼如下：
os.environ["CUDA_VISIBLE_DEVICES"] = "2"
補充知識：TensorFlow 設置程序可見GPU與邏輯分區
TensorFlow 設置程序可見GPU(多GPU情況)
|
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
|
import matplotlib as mplimport matplotlib.pyplot as plt%matplotlib inlineimport numpy as npimport sklearnimport pandas as pdimport osimport sysimport timeimport tensorflow as tffrom tensorflow_core.python.keras.api._v2 import kerasprint(tf.__version__)print(sys.version_info)for module in mpl, np, pd, sklearn, tf, keras: print(module.__name__, module.__version__)# 打印變量所在位置tf.debugging.set_log_device_placement(True) # 獲取物理GPU的個數gpus = tf.config.experimental.list_physical_devices("GPU") if len(gpus) >= 1: # 設置第幾個GPU 當前程序可見 tf.config.experimental.set_visible_devices(gpus[0], "GPU") print("物理GPU個數:", len(gpus))# 獲取邏輯GPU的個數logical_gpus = tf.config.experimental.list_logical_devices("GPU") print("邏輯GPU個數:", len(logical_gpus)) |
TensorFlow 設置GPU的 邏輯分區
|
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
|
import matplotlib as mplimport matplotlib.pyplot as plt%matplotlib inlineimport numpy as npimport sklearnimport pandas as pdimport osimport sysimport timeimport tensorflow as tffrom tensorflow_core.python.keras.api._v2 import kerasprint(tf.__version__)print(sys.version_info)for module in mpl, np, pd, sklearn, tf, keras: print(module.__name__, module.__version__)# 打印變量所在位置tf.debugging.set_log_device_placement(True) # 獲取物理GPU的個數gpus = tf.config.experimental.list_physical_devices("GPU") if len(gpus) >= 1: # 設置第幾個GPU 當前程序可見 tf.config.experimental.set_visible_devices(gpus[0], "GPU") # 設置GPU的 邏輯分區 tf.config.experimental.set_virtual_device_configuration( gpus[0], [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=3072), tf.config.experimental.VirtualDeviceConfiguration(memory_limit=3072)])print("物理GPU個數:", len(gpus))# 獲取邏輯GPU的個數logical_gpus = tf.config.experimental.list_logical_devices("GPU") print("邏輯GPU個數:", len(logical_gpus)) |
TensorFlow 手動設置處理GPU
|
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
|
import matplotlib as mplimport matplotlib.pyplot as plt%matplotlib inlineimport numpy as npimport sklearnimport pandas as pdimport osimport sysimport timeimport tensorflow as tffrom tensorflow_core.python.keras.api._v2 import kerasprint(tf.__version__)print(sys.version_info)for module in mpl, np, pd, sklearn, tf, keras: print(module.__name__, module.__version__)# 打印變量所在位置tf.debugging.set_log_device_placement(True) # 自動指定處理設備tf.config.set_soft_device_placement(True)# 獲取物理GPU的個數gpus = tf.config.experimental.list_physical_devices("GPU") for gpu in gpus: # 設置內存自增長方式 tf.config.experimental.set_memory_growth(gpu, True) print("物理GPU個數:", len(gpus))# 獲取邏輯GPU的個數logical_gpus = tf.config.experimental.list_logical_devices("GPU") print("邏輯GPU個數:", len(logical_gpus))c = []# 循環遍歷當前邏輯GPUfor gpu in logical_gpus: print(gpu.name) # 手動設置處理GPU with tf.device(gpu.name): a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]) # 矩陣相乘 并且添加至列表 c.append(tf.matmul(a, b))# 手動設置處理GPUwith tf.device("/GPU:0"): matmul_sum = tf.add_n(c)print(matmul_sum) |
以上這篇Tensorflow全局設置可見GPU編號操作就是小編分享給大家的全部內容了，希望能給大家一個參考，也希望大家多多支持服務器之家。
原文鏈接：https://blog.csdn.net/silent56_th/article/details/81319597
