tensorflow和tensorflow2.0控制显存
下面的方法可以控制tensorflow或keras实现显存自适应。
if tf.__version__.startswith('1.'): # tensorflow 1
config = tf.ConfigProto() # allow_soft_placement=True
config.gpu_options.allow_growth = True
sess = tf.Session(config=config)
else: # tensorflow 2
tf.config.gpu.set_per_process_memory_growth(enabled=True)
第一个方法用于控制tensorflow 1.x版本使用自适应显存,避免显存独占。第二个方法用于控制tensorflow 2.x 使用自适应显存。
原文地址:https://doit-space.blog.csdn.net/article/details/102911328