diff --git a/python/rapidocr_onnxruntime/config.yaml b/python/rapidocr_onnxruntime/config.yaml index 3bd206126..a8da9cc5e 100644 --- a/python/rapidocr_onnxruntime/config.yaml +++ b/python/rapidocr_onnxruntime/config.yaml @@ -7,7 +7,13 @@ Global: min_height: 30 width_height_ratio: 8 + intra_op_num_threads: &intra_nums -1 + inter_op_num_threads: &inter_nums -1 + Det: + intra_op_num_threads: *intra_nums + inter_op_num_threads: *inter_nums + use_cuda: false model_path: models/ch_PP-OCRv4_det_infer.onnx @@ -23,6 +29,9 @@ Det: score_mode: fast Cls: + intra_op_num_threads: *intra_nums + inter_op_num_threads: *inter_nums + use_cuda: false model_path: models/ch_ppocr_mobile_v2.0_cls_infer.onnx @@ -33,6 +42,9 @@ Cls: label_list: ['0', '180'] Rec: + intra_op_num_threads: *intra_nums + inter_op_num_threads: *inter_nums + use_cuda: false model_path: models/ch_PP-OCRv4_rec_infer.onnx diff --git a/python/rapidocr_onnxruntime/utils.py b/python/rapidocr_onnxruntime/utils.py index 2e6efa8f4..778066dd7 100644 --- a/python/rapidocr_onnxruntime/utils.py +++ b/python/rapidocr_onnxruntime/utils.py @@ -3,6 +3,7 @@ # @Contact: liekkaskono@163.com import argparse import math +import os import random import traceback import warnings @@ -33,6 +34,15 @@ def __init__(self, config): sess_opt.enable_cpu_mem_arena = False sess_opt.graph_optimization_level = GraphOptimizationLevel.ORT_ENABLE_ALL + cpu_nums = os.cpu_count() + intra_op_num_threads = config.get("intra_op_num_threads", -1) + if intra_op_num_threads != -1 and 1 <= intra_op_num_threads <= cpu_nums: + sess_opt.intra_op_num_threads = intra_op_num_threads + + inter_op_num_threads = config.get("inter_op_num_threads", -1) + if inter_op_num_threads != -1 and 1 <= inter_op_num_threads <= cpu_nums: + sess_opt.inter_op_num_threads = inter_op_num_threads + cpu_ep = "CPUExecutionProvider" cpu_provider_options = { "arena_extend_strategy": "kSameAsRequested",