当前位置:  开发笔记 > 编程语言 > 正文

tensorflow指定GPU与动态分配GPU memory设置

今天小编就为大家分享一篇tensorflow指定GPU与动态分配GPUmemory设置,具有很好的参考价值,希望对大家有所帮助。一起跟随小编过来看看吧

在tensorflow中,默认指定占用所有的GPU,如需指定占用的GPU,可以在命令行中:

export CUDA_VISIBLE_DEVICES=1

这样便是只占用1号GPU,通过命令

nvidia-smi

可以查看各个GPU的使用情况。

另外,也可以在python程序中指定GPU,并且动态分配memory,代码如下

import os
import sys
os.environ['CUDA_VISIBLE_DEVICES'] = sys.argv[1]
import tensorflow as tf
from keras.backend.tensorflow_backend import set_session
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
set_session(tf.Session(config=config))

这样,占用的GPU就是sys.argv[1]了,并且会按需分配memory。

以上这篇tensorflow指定GPU与动态分配GPU memory设置就是小编分享给大家的全部内容了,希望能给大家一个参考,也希望大家多多支持。

推荐阅读
雯颜哥_135
这个屌丝很懒,什么也没留下!
DevBox开发工具箱 | 专业的在线开发工具网站    京公网安备 11010802040832号  |  京ICP备19059560号-6
Copyright © 1998 - 2020 DevBox.CN. All Rights Reserved devBox.cn 开发工具箱 版权所有