tensorflow显存占用控制

0x01 问题描述

Tensorflow使用时默认占满所有显卡所有的显存,在公用计算机上导致一个用户使用tensorflow,其他用户没法使用.尤其是使用jupyter notebook时更甚,因为用户常常忘记关进程,导致GPU显存资源不能释放,出现GPU使用率为0,但显存占满的情况.

0x02 解决方案

# 设置可见GPU
gpu_no = '0' # or '1'
os.environ["CUDA_VISIBLE_DEVICES"] = gpu_no
#定义TensorFlow配置
config = tf.ConfigProto()
#配置GPU内存分配方式
config.gpu_options.allow_growth = True
config.gpu_options.per_process_gpu_memory_fraction = 0.4
with tf.Session(config=config, ...)) as sess:
    pass

0x03 对比实验


原创文章,转载请注明: 转载自张哲的博客

本文链接地址: tensorflow显存占用控制