gpu_utils.py 文件源码

def get_n_gpu():
    detected_n_gpu = mp.RawValue('i', 0)
    p = mp.Process(target=n_gpu_subprocess, args=(detected_n_gpu,))
    p.start()
    p.join()
    n_gpu = int(detected_n_gpu.value)
    if n_gpu == -1:
        raise ImportError("Must be able to import pygpu to use GPUs.")
    return n_gpu