設置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的使用與分配