配置动态显存的必要性
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进行配置。