用tensorflow的同学肯定都知道,只要tf的程序运行起来就会默认占满显存,但实际上有些代码使用的数据量远远不需要使用全部显存,这样就造成了浪费,我们可以在代码前加上一些控制语句,使程序占用的显存智能一些。
指定显卡
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
对于Tensorflow1.x版本
# 按比例配置显存
config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.4 # 根据自己的需求确定
session = tf.Session(config=config, ...)
# 自适应显存占用
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
session = tf.Session(config=config, ...)
对于Tensorflow2.x版本
# 自适应显存占用
gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
for gpu in gpus:
tf.config.experimental.set_memory_growth(gpu, True)