exct.py 文件源码

 复制代码def init_gpus(rank, n_parallel=None):
    import theano
    import theano.gpuarray
    try:
        theano.gpuarray.use("cuda" + str(rank))
    except Exception as exc:
        if n_parallel is not None:
            raise exc("Master unable to use GPU.")
        else:
            sync.workers_OK.value = False
            raise exc("Worker rank {} unable to use GPU.".format(rank))
    finally:
        sync.barrier_out.wait()
    if n_parallel is not None:
        if sync.workers_OK.value:
            print("Synkhronos: {} GPUs initialized, master rank: {}".format(
                n_parallel, rank))
        else:
            raise RuntimeError("Workers did not initialize GPUs.")