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

笔者需要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)

来源：https://blog.csdn.net/silent56_th/article/details/81319597