add cpu_math_library_num_threads params

Author: LDOUBLEV

tools/infer/utility.py

@@ -21,6 +21,9 @@
 from PIL import Image, ImageDraw, ImageFont
 import math
 from paddle import inference
+import time
+from ppocr.utils.logging import get_logger
+logger = get_logger()
 
 
 def parse_args():
@@ -98,6 +101,7 @@ def str2bool(v):
     parser.add_argument("--cls_thresh", type=float, default=0.9)
 
     parser.add_argument("--enable_mkldnn", type=str2bool, default=False)
+    parser.add_argument("--cpu_threads", type=int, default=10)
     parser.add_argument("--use_pdserving", type=str2bool, default=False)
 
     parser.add_argument("--use_mp", type=str2bool, default=False)
@@ -140,14 +144,15 @@ def create_predictor(args, mode, logger):
                 max_batch_size=args.max_batch_size)
     else:
         config.disable_gpu()
-        config.set_cpu_math_library_num_threads(6)
+        if hasattr(args, "cpu_threads"):
+            config.set_cpu_math_library_num_threads(args.cpu_threads)
+        else:
+            config.set_cpu_math_library_num_threads(
+                10)  # default cpu threads as 10
         if args.enable_mkldnn:
             # cache 10 different shapes for mkldnn to avoid memory leak
             config.set_mkldnn_cache_capacity(10)
             config.enable_mkldnn()
-            #  TODO LDOUBLEV: fix mkldnn bug when bach_size  > 1
-            #config.set_mkldnn_op({'conv2d', 'depthwise_conv2d', 'pool2d', 'batch_norm'})
-            args.rec_batch_num = 1
 
     # enable memory optim
     config.enable_memory_optim()