config.py 文件源码

def _query_gpu_info():
  """ This function query GPU information:
  ngpu
  [device_name, device_compute_capability, device_total_memory]

  Note
  ----
  this function use deviceQuery command, so you better have it
  in your path
  """
  dev = {'ngpu': 1,
     # deviceName: [cardName, computeCapability, mem(MB)]
     'dev0': ['Unknown', 3.0, 1024]}
  temp_dir = tempfile.mkdtemp()
  p = os.path.join(temp_dir, 'tmp.txt')
  queried = subprocess.call('deviceQuery > ' + p,
          shell=True,
          stdout=subprocess.PIPE,
          stderr=subprocess.PIPE) == 0
  dev = {}
  if queried: # found deviceQuery
    info = open(p, 'r').read()
    devNames = re.compile(r'Device \d: ".*"').findall(info)
    devNames = [i.strip().split(':')[-1].replace('"', '') for i in devNames]
    ngpu = len(devNames)
    comCap = re.compile(
        r'CUDA Capability Major\/Minor version number:\s*.*').findall(info)
    comCap = [float(i.strip().split(':')[-1]) for i in comCap]
    totalMems = re.compile(
        r'Total amount of global memory:\s*\d*').findall(info)
    totalMems = [int(i.strip().split(':')[-1]) for i in totalMems]
    # ====== create dev ====== #
    dev['ngpu'] = ngpu
    for i, (name, com, mem) in enumerate(zip(devNames, comCap, totalMems)):
      dev['dev%d' % i] = [name, com, mem]
  else:
    _warning('Cannot use "deviceQuery" to get GPU information for configuration.')
    from tensorflow.python.client import device_lib
    local_device_protos = device_lib.list_local_devices()
    dev['ngpu'] = 0
    for i, name in (x.name for x in local_device_protos
        if x.device_type == 'GPU'):
      dev['dev%d' % i] = [name, None, None]
      dev['ngpu'] += 1
  # remove temp-dir
  shutil.rmtree(temp_dir)
  return dev