tensorflow和tensorflow2.0控制显存

Doit_行之

于 2019-11-05 10:41:09 发布

阅读量6.7k

点赞数 6

CC 4.0 BY-SA版权

分类专栏： TensorFlow2教程文章标签： tensorflow 显存

本文链接：https://blog.csdn.net/qq_31456593/article/details/102911328

TensorFlow2教程专栏收录该内容

40 篇文章

订阅专栏

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