From d66d6c6ea355832243667ea5a01add40fb3e8f73 Mon Sep 17 00:00:00 2001 From: tensor-tang Date: Wed, 15 Nov 2017 17:21:27 +0800 Subject: [PATCH] auto set cpu environment in V2 API --- python/paddle/v2/__init__.py | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/python/paddle/v2/__init__.py b/python/paddle/v2/__init__.py index 3d705138434..a55b9d7a21b 100644 --- a/python/paddle/v2/__init__.py +++ b/python/paddle/v2/__init__.py @@ -78,6 +78,31 @@ def init(**kwargs): for key in args_dict.keys(): args.append('--%s=%s' % (key, str(args_dict[key]))) + # auto set cpu environment + def set_env(key, value): + '''If the key has not been set in the environment, set it with value.''' + assert isinstance(key, str) + assert isinstance(value, str) + envset = os.environ.get(key) + if envset is None: + os.environ[key] = value + + ht = os.popen("lscpu |grep \"per core\"|awk -F':' '{print $2}'|xargs") + ht = int(ht.read()) + if ht == 1: # ht is off + set_env("OMP_DYNAMIC", "false") + set_env("KMP_AFFINITY", "granularity=fine,compact,0,0") + else: + set_env("OMP_DYNAMIC", "true") + set_env("KMP_AFFINITY", "granularity=fine,compact,1,0") + processors = os.popen("grep \"processor\" /proc/cpuinfo|sort -u|wc -l") + processors = int(processors.read()) + trainers = kwargs.get('trainer_count', 1) + threads = processors / trainers + threads = '1' if threads < 1 else str(threads) + set_env("OMP_NUM_THREADS", threads) + set_env("MKL_NUM_THREADS", threads) + if 'use_gpu' in kwargs: cp.g_command_config_args['use_gpu'] = kwargs['use_gpu'] if 'use_mkldnn' in kwargs: -- GitLab