默认情况下,TensorFlow 将使用几乎所有可用的显存,以避免内存碎片化所带来的性能损失,但这样不能在一台机器上运行多个程序
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)))
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)
原文:https://www.cnblogs.com/yu212223/p/12145130.html