From 77c80e48da393abf81ff037478bcd57e29e9ea2a Mon Sep 17 00:00:00 2001 From: lizhenyu Date: Wed, 30 Dec 2020 10:46:16 +0800 Subject: [PATCH] bugfix: ps cache data process performance decay --- mindspore/parallel/_ps_context.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/mindspore/parallel/_ps_context.py b/mindspore/parallel/_ps_context.py index 40b4ad6266..b14c8c8f2d 100644 --- a/mindspore/parallel/_ps_context.py +++ b/mindspore/parallel/_ps_context.py @@ -14,6 +14,7 @@ # ============================================================================ """Context for parameter server training mode""" +import os from mindspore._c_expression import PSContext _ps_context = None @@ -134,4 +135,10 @@ def _clone_hash_table(dest_param_name, src_param_name): ps_context().clone_hash_table(dest_param_name, src_param_name) def _set_cache_enable(cache_enable): + # Environment variables are used to specify a maximum number of OpenBLAS threads: + # In ubuntu(GPU) environment, numpy will use too many threads for computing, + if cache_enable: + os.environ['OPENBLAS_NUM_THREADS'] = '2' + os.environ['GOTO_NUM_THREADS'] = '2' + os.environ['OMP_NUM_THREADS'] = '2' ps_context().set_cache_enable(cache_enable)