设置tensorflow的显存为动态使用

设置tensorflow的显存为动态使用

默认情况下，TensorFlow 将使用几乎所有可用的显存，以避免内存碎片化所带来的性能损失，但这样不能在一台机器上运行多个程序

tensorflow 1.x

from tensorflow.compat.v1 import GPUOptions
from tensorflow.compat.v1 import ConfigProto
from tensorflow.compat.v1 import Session
from keras.backend.tensorflow_backend import set_session

gpu_options = GPUOptions(allow_growth=True)
set_session(Session(config=ConfigProto(gpu_options=gpu_options)))

详见tensorflow入门笔记1:指定GPU及分配显存

tensorflow 2.0

import tensorflow as tf

gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)

详见Tensorflow 2.0 GPU的使用与分配