项目中需要控制tensorflow的显存占用,如何控制gpu显存的方法网上有很多,但都不起作用,调试之后发现如下代码问题:
from tensorflow.python.client import device_lib
devices = device_lib.list_local_devices()
一旦调用了上面函数,从log看就会createdevice,而且将显存占满。将上面的函数剃掉,在控制设备运行环境时直接如下:
with tf.device('/gpu:0'), tf.name_scope('%s_%d' % (GPU_NAME, 0)) as scope:
然后再使用网上控制gpu显存的方法,就可以了。