tensorflow显存自适应分配(tf.1x和tf2.x)

用tensorflow的同学肯定都知道,只要tf的程序运行起来就会默认占满显存,但实际上有些代码使用的数据量远远不需要使用全部显存,这样就造成了浪费,我们可以在代码前加上一些控制语句,使程序占用的显存智能一些。

# 指定显卡
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
......
# tf1.x version  
# # 按比例配置显存
config = tf.ConfigProto()

config.gpu_options.per_process_gpu_memory_fraction = 0.7 # 根据自己的需求确定
# 上面一行还可以改为下面两行
config = tf.ConfigProto(allow_soft_placement=True)
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.7)  #o.7可以自己修改其他值

with tf.Session(config=config) as sess:

# tf1.x version  
# 自适应显存占用
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
with tf.Session(config=config) as sess:



# tf2.x version  自适应显存占用
# gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
# for gpu in gpus:
#     tf.config.experimental.set_memory_growth(gpu, True)

猜你喜欢

转载自blog.csdn.net/SMF0504/article/details/106850769
今日推荐