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