add cpu_math_library_num_threads params
@@ -21,6 +21,9 @@ import json
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 parse_args():
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()
