keras资源设置

Posted on 2017-09-13(星期三) 14:00 in Data

  • 1.关闭GPU,只使用CPU

    在import tensorflow之前,加上:

    import os
    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"   # see issue #152
    os.environ["CUDA_VISIBLE_DEVICES"] = ""
    
  • 2.设置keras占用GPU内存的比例:

    import os
    import tensorflow as tf
    import keras.backend.tensorflow_backend as K
    
    def get_session(gpu_fraction=0.3):
        '''Assume that you have 6GB of GPU memory and want to allocate ~2GB'''
    
        num_threads = os.environ.get('OMP_NUM_THREADS')
        gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=gpu_fraction)
    
        if num_threads:
            return tf.Session(config=tf.ConfigProto(
                gpu_options=gpu_options, intra_op_parallelism_threads=num_threads))
        else:
            return tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))
    
    K.set_session(get_session())
    

    两种限定GPU占用的方法:

    (1)在tensorflow中定义session时作如下设置,该设置会启用最少的GPU显存来运行程序。

    config = tf.ConfigProto() 
    config.gpu_options.allow_growth = True 
    session = tf.Session(config=config)
    

    (2)在tensorflow中定义session时作如下设置,该设置会强制程序只占用指定比例的GPU显存。

    config = tf.ConfigProto() 
    config.gpu_options.per_process_gpu_memory_fraction = 0.4 # 占用GPU40%的显存 
    session = tf.Session(config=config)
    
  • 3 设置keras分配指定GPU:

    通过 watch -n 1 nvidia-smi 命令查看 GPU 信息

    设置使用GPU序号为‘0’的GPU(在set_session之前):

    os.environ["CUDA_VISIBLE_DEVICES"] = '0'