|
|
|
@ -82,7 +82,24 @@ def _current_expected_place():
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def _cpu_num():
|
|
|
|
|
return int(os.environ.get('CPU_NUM', multiprocessing.cpu_count()))
|
|
|
|
|
cpu_num = os.environ.get('CPU_NUM', None)
|
|
|
|
|
if cpu_num is None:
|
|
|
|
|
sys.stderr.write(
|
|
|
|
|
'The CPU_NUM is not specified, you should set CPU_NUM in '
|
|
|
|
|
'the environment variable list, i.e export CPU_NUM=1. CPU_NUM '
|
|
|
|
|
'indicates that how many CPUPlace are used in the current task.\n'
|
|
|
|
|
'!!! The default number of CPUPlaces is 1.')
|
|
|
|
|
os.environ['CPU_NUM'] = str(1)
|
|
|
|
|
return int(cpu_num)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def _cuda_ids():
|
|
|
|
|
gpus_env = os.getenv("FLAGS_selected_gpus")
|
|
|
|
|
if gpus_env:
|
|
|
|
|
device_ids = [int(s) for s in gpus_env.split(",")]
|
|
|
|
|
else:
|
|
|
|
|
device_ids = six.moves.range(core.get_cuda_device_count())
|
|
|
|
|
return device_ids
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def cuda_places(device_ids=None):
|
|
|
|
@ -116,11 +133,7 @@ def cuda_places(device_ids=None):
|
|
|
|
|
assert core.is_compiled_with_cuda(), \
|
|
|
|
|
"Not compiled with CUDA"
|
|
|
|
|
if device_ids is None:
|
|
|
|
|
gpus_env = os.getenv("FLAGS_selected_gpus")
|
|
|
|
|
if gpus_env:
|
|
|
|
|
device_ids = [int(s) for s in gpus_env.split(",")]
|
|
|
|
|
else:
|
|
|
|
|
device_ids = six.moves.range(core.get_cuda_device_count())
|
|
|
|
|
device_ids = _cuda_ids()
|
|
|
|
|
elif not isinstance(device_ids, (list, tuple)):
|
|
|
|
|
device_ids = [device_ids]
|
|
|
|
|
return [core.CUDAPlace(dev_id) for dev_id in device_ids]
|
|
|
|
|