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

笔者需要tensorflow仅运行在一个GPU上（机器本身有多GPU），而且需要依据系统参数动态调节，故无法简单使用CUDA_VISIBLE_DEVICES。

一种方式是全局使用tf.device函数生成的域，但设备号需要在绘制Graph前指定，仍然不够灵活。

查阅文档发现config的GPUOptions中的visible_device_list可以定义GPU编号从visible到virtual的映射，即可以设置tensorflow可见的GPU device，从而全局设置了tensorflow可见的GPU编号。代码如下：

									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情况)

									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的逻辑分区

									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

									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)