Tensorflow全局设置可见GPU编号操作
笔者需要tensorflow仅运行在一个GPU上（机器本身有多GPU），而且需要依据系统参数动态调节，故无法简单使用CUDA_VISIBLE_DEVICES。
一种方式是全局使用tf.device函数生成的域，但设备号需要在绘制Graph前指定，仍然不够灵活。
查阅文档发现config的GPUOptions中的visible_device_list可以定义GPU编号从visible到virtual的映射，即可以设置tensorflow可见的GPUdevice，从而全局设置了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情况)
importmatplotlibasmpl
importmatplotlib.pyplotasplt
%matplotlibinline
importnumpyasnp
importsklearn
importpandasaspd
importos
importsys
importtime
importtensorflowastf
fromtensorflow_core.python.keras.api._v2importkeras
print(tf.__version__)
print(sys.version_info)
formoduleinmpl,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")
iflen(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的逻辑分区
importmatplotlibasmpl
importmatplotlib.pyplotasplt
%matplotlibinline
importnumpyasnp
importsklearn
importpandasaspd
importos
importsys
importtime
importtensorflowastf
fromtensorflow_core.python.keras.api._v2importkeras
print(tf.__version__)
print(sys.version_info)
formoduleinmpl,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")
iflen(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
importmatplotlibasmpl
importmatplotlib.pyplotasplt
%matplotlibinline
importnumpyasnp
importsklearn
importpandasaspd
importos
importsys
importtime
importtensorflowastf
fromtensorflow_core.python.keras.api._v2importkeras
print(tf.__version__)
print(sys.version_info)
formoduleinmpl,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")
forgpuingpus:
#设置内存自增长方式
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
forgpuinlogical_gpus:
print(gpu.name)
#手动设置处理GPU
withtf.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
withtf.device("/GPU:0"):
matmul_sum=tf.add_n(c)
print(matmul_sum)
以上这篇Tensorflow全局设置可见GPU编号操作就是小编分享给大家的全部内容了，希望能给大家一个参考，也希望大家多多支持毛票票。