Tensorflow全局设置可见GPU编号操作

时间:2022-03-20 18:55:48

笔者需要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 mpl
import matplotlib.pyplot as plt
%matplotlib inline
import numpy as np
import sklearn
import pandas as pd
import os
import sys
import time
import tensorflow as tf
 
from tensorflow_core.python.keras.api._v2 import keras
 
print(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 mpl
import matplotlib.pyplot as plt
%matplotlib inline
import numpy as np
import sklearn
import pandas as pd
import os
import sys
import time
import tensorflow as tf
 
from tensorflow_core.python.keras.api._v2 import keras
 
print(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 mpl
import matplotlib.pyplot as plt
%matplotlib inline
import numpy as np
import sklearn
import pandas as pd
import os
import sys
import time
import tensorflow as tf
 
from tensorflow_core.python.keras.api._v2 import keras
 
print(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 = []
 
# 循环遍历当前逻辑GPU
for 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))
 
# 手动设置处理GPU
with tf.device("/GPU:0"):
 matmul_sum = tf.add_n(c)
 
print(matmul_sum)

以上这篇Tensorflow全局设置可见GPU编号操作就是小编分享给大家的全部内容了,希望能给大家一个参考,也希望大家多多支持服务器之家。

原文链接:https://blog.csdn.net/silent56_th/article/details/81319597