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)