import tensorflow as tf
import tensorflow.keras.backend as K
os.environ["TF_FORCE_GPU_ALLOW_GROWTH"] = "true"
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.Session(config=config)
K.set_session(sess)
加入上面几行,一定程度上能解决显存溢出的问题。
我用的是Tensorflow1.15.0-gpu