您的位置:首页 > 其它

TensorFlow设置GPU占用量

2018-08-19 16:36 316 查看

tensorflow在训练时默认占用所有GPU的显存。

可以通过以下方式解决该问题:

1、在构造tf.Session()时候通过传递tf.GPUOptions作为可选配置参数的一部分来显式地指定需要分配的显存比例,如下所示:

# 假如有12GB的显存并使用其中的4GB:
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.333)
sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))

per_process_gpu_memory_fraction指定了每个GPU进程中使用显存的上限,但它只能均匀作用于所有GPU,无法对不同GPU设置不同的上限

2、尝试如下设置:

config = tf.ConfigProto()
config.gpu_options.allow_growth=True
sess = tf.Session(config=config)

当allow_growth设置为True时,分配器将不会指定所有的GPU内存,而是根据需求增长

3、指定GPU编号:

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

或者在脚本或者命令行中指定
export CUDA_VISIBLE_DEVICES=0

 

内容来自用户分享和网络整理,不保证内容的准确性,如有侵权内容,可联系管理员处理 点击这里给我发消息
标签: