sg_main.py 文件源码

def sg_gpus():
    r""" Gets current available GPU nums

    Returns:
      A integer : total # of GPUs available
    """
    global _gpus

    if _gpus is None:
        local_device_protos = device_lib.list_local_devices()
        _gpus = len([x.name for x in local_device_protos if x.device_type == 'GPU'])

    return max(_gpus, 1)


#
# context helpers
#