因为涉及到同一台电脑多个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