tensorflow和tensorflow2.0控制顯存

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