tensorflow 导入新的tensorflow实例

时间:2023-03-09 13:19:43
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'] = '' if verbose else ''
import tensorflow as tf
tf.logging.set_verbosity(tf.logging.DEBUG if verbose else tf.logging.ERROR)
return tf