tensorflow--指定GPU及分配显存

最新推荐文章于 2024-05-20 15:40:11 发布

转载最新推荐文章于 2024-05-20 15:40:11 发布 · 468 阅读

1 ·

CC 4.0 BY-SA版权

原文链接：https://www.jianshu.com/p/756541f9ff40

tensorflow，pytorch 专栏收录该内容

9 篇文章

订阅专栏

export CUDA_VISIBLE_DEVICES]

转载：https://www.jianshu.com/p/756541f9ff40
前提：
1.安装好tensorflow的gpu版本及其对应CUDA
2.有GPU

指定某一块或多块gpu运行

1

# 指定第二块gpu
CUDA_VISIBLE_DEVICES=1 python {xxx}.py

2

export CUDA_VISIBLE_DEVICES=1
python {xxx}.py

3

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "2"

如果设置-1，是不使用显卡。
参数参考：

Environment Variable Syntax	Results
CUDA_VISIBLE_DEVICES=1	CUDA_VISIBLE_DEVICES=1
CUDA_VISIBLE_DEVICES=0,1	Devices 0 and 1 will be visible
CUDA_VISIBLE_DEVICES=“0,1”	Devices 0 and 1 will be visible
CUDA_VISIBLE_DEVICES=“0,1”	Devices 0, 2, 3 will be visible; device 1 is masked
CUDA_VISIBLE_DEVICES=""	No GPU will be visible

配置gpu显存使用率

方法一：通过配置Session的运行参数配置gpu的使用

# 通过配置Session的运行参数配置gpu的使用

config = tf.ConfigProto(allow_soft_placement=True, allow_soft_placement=True)
config.gpu_options.per_process_gpu_memory_fraction = 0.4  #占用40%显存
sess = tf.Session(config=config）

参数参考：

Syntax	Results
tf.ConfigProto(log_device_placement=True)	记录设备指派情况
tf.ConfigProto(allow_soft_placement=True）	自动选择一个运行设备
tf.ConfigProto(allow_soft_placement=True）	动态申请显存
config.gpu_options.per_process_gpu_memory_fraction = 0.4	动态申请显存