Add 'kmp-blocktime' for mlperf-gnmt (#494)

* Add 'kmp-blocktime' for mlperf-gnmt Signed-off-by: Abolfazl Shahbazi <abolfazl.shahbazi@intel.com> * Remove duplicate parameter definition Signed-off-by: Abolfazl Shahbazi <abolfazl.shahbazi@intel.com>
ZhaoqiongZ · Apr 5, 2022 · b022eef · b022eef
1 parent 2751157
commit b022eef
Showing 1 changed file with 7 additions and 4 deletions.
diff --git a/benchmarks/language_translation/tensorflow/mlperf_gnmt/inference/fp32/model_init.py b/benchmarks/language_translation/tensorflow/mlperf_gnmt/inference/fp32/model_init.py
@@ -36,10 +36,6 @@ def __init__(self, args, custom_args=[], platform_util=None):
                         (str(self.args.num_cores - 1)) + " "
         self.cmd += "{} ".format(self.python_exe)
 
-        # Set KMP env vars, if they haven't already been set
-        config_file_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "config.json")
-        self.set_kmp_vars(config_file_path)
-
         # use default batch size if -1
         if self.args.batch_size == -1:
             self.args.batch_size = 32
@@ -48,7 +44,14 @@ def __init__(self, args, custom_args=[], platform_util=None):
         self.set_num_inter_intra_threads()
         set_env_var("OMP_NUM_THREADS", self.args.num_intra_threads)
         arg_parser = ArgumentParser(description="process custom_args")
+        arg_parser.add_argument('--kmp-blocktime', dest='kmp_blocktime',
+                                help='number of kmp block time',
+                                type=int, default=1)
         self.args = arg_parser.parse_args(self.custom_args, namespace=self.args)
+        # Set KMP env vars, if they haven't already been set, but override the default KMP_BLOCKTIME value
+        config_file_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "config.json")
+        self.set_kmp_vars(config_file_path, kmp_blocktime=str(self.args.kmp_blocktime))
+
         src_vocab_file = os.path.join(self.args.data_location, "vocab.bpe.32000.en")
         tgt_vocab_file = os.path.join(self.args.data_location, "vocab.bpe.32000.de")
         inference_input_file = os.path.join(self.args.data_location, "newstest2014.tok.bpe.32000.en")