tensorflow 导入新的tensorflow实例

因为涉及到同一台电脑多个GPU,在指定tensorflow图的时候,需要为不同的图指定不同的GPU,所以必须在导入tensorflow之前,指定可用的GPU

def import_tf(device_id=-1, verbose=False):
    """
    import tensorflow, set tensorflow graph load device, set tensorflow log level, return tensorflow instance
    :param device_id: GPU id
    :param verbose: tensorflow logging level
    :return: tensorflow instance
    """
    # set visible gpu, -1 is cpu
    os.environ['CUDA_VISIBLE_DEVICES'] = '-1' if device_id < 0 else str(device_id)
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '0' if verbose else '3'
    import tensorflow as tf
    tf.logging.set_verbosity(tf.logging.DEBUG if verbose else tf.logging.ERROR)
    return tf

 

posted @ 2019-02-16 17:26  下路派出所  阅读(353)  评论(0编辑  收藏  举报