From 16d70fb74bbeded98c134df0ab69bfc4b29eff31 Mon Sep 17 00:00:00 2001 From: LDOUBLEV Date: Wed, 26 May 2021 07:08:56 +0000 Subject: [PATCH] add cpu_math_library_num_threads params --- tools/infer/utility.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/tools/infer/utility.py b/tools/infer/utility.py index b5fe3ba9..41ca445f 100755 --- a/tools/infer/utility.py +++ b/tools/infer/utility.py @@ -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() -- GitLab