TensorFlow1.x/2.x配置动态显存方法

配置动态显存的必要性

TensorFlow默认是占用全部的显存，不过比较BUG的是，就算模型用不了全部的显存，TensorFlow也会尽可能地多申请显存，可能是出于性能方面的考虑，并且如果满足不了它就会罢工。。比如我一块6GB显存的1660Ti或者RTX 2060，运行YOLOv4的推理程序如果不配置显存动态增长就会报错，不管是TensorFlow 1.x还是2.x的版本我都遇到了这样的问题。明明用darknet运行啥事都没有，就这货会搞事情。。下面简单记录下TensorFlow 1.x / 2.x版本的配置动态显存增长的方法。

配置方法

对于TensorFlow 1.x版本的用户，配置的方法如下：

import tensorflow as tf

config = tf.ConfigProto(allow_soft_placement=True)
config.gpu_options.allow_growth = True
with tf.Session(config=config) as sess:
	# sess.run(your tasks)

而对于TensorFlow 2.x版本，配置方法则有所不同。以最新的2.3.0版本为例，配置的方法如下：

import tensorflow as tf

tf.config.set_soft_device_placement(True)
physical_devices = tf.config.list_physical_devices('GPU')
if len(physical_devices) > 0:
	tf.config.experimental.set_memory_growth(physical_devices[0], True)

# run your tasks

需要注意的是，在2.x版本中要把配置的代码写在调用TensorFlow API之前，不然系统会报错。然后，上面的代码是对单个GPU进行的配置，有多个GPU的可以分别指定device进行配置。