main.py 文件源码

def restrict_gpu_memory(per_process_gpu_memory_fraction: float = 0.9):
    import os
    import tensorflow as tf
    import keras
    thread_count = os.environ.get('OMP_NUM_THREADS')
    gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=per_process_gpu_memory_fraction)
    config = tf.ConfigProto(gpu_options=gpu_options,
                            allow_soft_placement=True,
                            intra_op_parallelism_threads=thread_count) \
        if thread_count else tf.ConfigProto(gpu_options=gpu_options, allow_soft_placement=True)
    keras.backend.tensorflow_backend.set_session(tf.Session(config=config))