TensorFlow全局设置可见GPU编号与显存管理指南
本文详解在TensorFlow中全局限制可见GPU编号的两种核心方法：通过ConfigProto配置visible_device_list以及直接设置环境变量CUDA_VISIBLE_DEVICES。针对TensorFlow 2.x，进一步介绍了利用tf.config.experimental API实现指定GPU可见、逻辑分区（VirtualDeviceConfiguration）及显存动态增长（memory_growth）的高级配置技巧，帮助开发者在多GPU环境中优化资源分配与模型部署效率。
TensorFlow全局限制可见GPU的两种基础方法
在多GPU服务器环境中，若需强制TensorFlow仅使用特定GPU，且该选择需依据系统参数动态调整，直接依赖环境变量CUDA_VISIBLE_DEVICES往往难以满足灵活性需求。
另一种常见做法是利用tf.device函数创建设备域，但此方式要求在构建计算图之前明确指定设备编号，对于动态场景而言仍显僵化。
查阅官方文档可知，通过配置tf.ConfigProto中的GPUOptions.visible_device_list属性，可以建立从可见物理设备到虚拟设备的映射关系，从而在全局层面控制TensorFlow可识别的GPU列表。具体实现代码如下：
config = tf.ConfigProto()
config.gpu_options.visible_device_list = str(device_num)
sess = tf.Session(config=config)此外，参考多卡服务器环境下隐藏部分GPU及显存设置的实践，也可以直接通过Python的os模块设置全局环境变量CUDA_VISIBLE_DEVICES来实现相同效果，代码示例如下：
os.environ["CUDA_VISIBLE_DEVICES"] = "2"
TensorFlow 2.x高级GPU配置：可见性、分区与显存管理
补充知识：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中手动指定计算设备的方法
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编号的主要操作，旨在为读者提供实用的参考。希望这些信息能帮助大家更好地管理GPU资源，同时也欢迎继续支持本站获取更多技术分享。
免责声明：本文内容仅供参考，不代表本站观点。转载或引用请注明出处。
