Tensorflow1.x及2.0指定显卡、自适应显存

用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)