Fix for Recognizing --threads Option in Benchmarking

vision/classification_and_detection/python/backend.py

@@ -16,7 +16,7 @@ def version(self):
     def name(self):
         raise NotImplementedError("Backend:name")
 
-    def load(self, model_path, inputs=None, outputs=None):
+    def load(self, model_path, inputs=None, outputs=None, threads=None):
         raise NotImplementedError("Backend:load")
 
     def predict(self, feed):

vision/classification_and_detection/python/backend_ncnn.py

@@ -18,7 +18,7 @@ def image_format(self):
         """image_format. For ncnn it is NCHW."""
         return "NCHW"
 
-    def load(self, model_path, inputs=None, outputs=None):
+    def load(self, model_path, inputs=None, outputs=None, threads=None):
         param_file, bin_file = f"{model_path}.param", f"{model_path}.bin"
         if param_file.endswith("resnet50_v1.param"):
             # download model files if doesn't

vision/classification_and_detection/python/backend_onnxruntime.py

@@ -24,9 +24,14 @@ def image_format(self):
         """image_format. For onnx it is always NCHW."""
         return "NCHW"
 
-    def load(self, model_path, inputs=None, outputs=None):
+    def load(self, model_path, inputs=None, outputs=None, threads=None):
         """Load model and find input/outputs from the model file."""
+        print("************************************************************")
+        print(">>> Value of num_threads: ", threads)
+        print("************************************************************")
         opt = rt.SessionOptions()
+        opt.intra_op_num_threads = int(threads)
+        opt.inter_op_num_threads = int(threads)
 
         # By default all optimizations are enabled
         # https://onnxruntime.ai/docs/performance/graph-optimizations.html

vision/classification_and_detection/python/backend_pytorch.py

@@ -30,8 +30,10 @@ def name(self):
     def image_format(self):
         return "NCHW"
 
-    def load(self, model_path, inputs=None, outputs=None):
+    def load(self, model_path, inputs=None, outputs=None, threads=None):
         self.model = onnx.load(model_path)
+        torch.set_num_threads(threads)
+        torch.set_num_interop_threads(threads)
 
         # find inputs from the model if not passed in by config
         if inputs:

vision/classification_and_detection/python/backend_pytorch_native.py

@@ -24,9 +24,11 @@ def name(self):
     def image_format(self):
         return "NCHW"
 
-    def load(self, model_path, inputs=None, outputs=None):
+    def load(self, model_path, inputs=None, outputs=None, threads=None):
         self.model = torch.load(model_path, weights_only=False)
         self.model.eval()
+        torch.set_num_threads(threads)
+        torch.set_num_interop_threads(threads)
         # find inputs from the model if not passed in by config
         if inputs:
             self.inputs = inputs

vision/classification_and_detection/python/backend_tf.py

@@ -27,7 +27,7 @@ def image_format(self):
         # NHWC)
         return "NHWC"
 
-    def load(self, model_path, inputs=None, outputs=None):
+    def load(self, model_path, inputs=None, outputs=None, threads=None):
         # there is no input/output meta data i the graph so it need to come
         # from config.
         if not inputs:

vision/classification_and_detection/python/backend_tvm.py

@@ -200,7 +200,7 @@ def _worker_handler(feed):
         global global_executor
         return global_executor.predict(feed)
 
-    def load(self, model_path, inputs=None, outputs=None):
+    def load(self, model_path, inputs=None, outputs=None, threads=None):
         """Load model and find input/outputs from the model file."""
         self.load_impl(model_path, inputs, outputs, self.max_batchsize)
 

vision/classification_and_detection/python/main.py

@@ -383,6 +383,18 @@ def get_args():
 
     if args.scenario not in SCENARIO_MAP:
         parser.error("valid scanarios:" + str(list(SCENARIO_MAP.keys())))
+
+    if args.threads:
+        num_threads = args.threads
+        print("*****>>> Value of num_threads: {num_threads} **********")
+        os.environ["OMP_NUM_THREADS"] = str(num_threads)
+        os.environ["OPENBLAS_NUM_THREADS"] = str(num_threads)
+        os.environ["MKL_NUM_THREADS"] = str(num_threads)
+        os.environ["TF_NUM_INTRAOP_THREADS"] = str(num_threads)
+        os.environ["TF_NUM_INTEROP_THREADS"] = str(num_threads)
+        os.environ["TVM_NUM_THREADS"] = str(num_threads)
+        os.environ["TFLITE_NUM_THREADS"] = str(num_threads)
+        os.environ["NCNN_NUM_THREADS"] = str(num_threads)
     return args
 
 
@@ -646,7 +658,9 @@ def main():
         model = backend.load(
             args.model,
             inputs=args.inputs,
-            outputs=args.outputs)
+            outputs=args.outputs,
+            threads=args.threads,
+        )
     final_results = {
         "runtime": model.name(),
         "version": model.version(),