diff --git a/deploy/hubserving/ocr_det/module.py b/deploy/hubserving/ocr_det/module.py index e5fac23faff29ff1b407fa5629b39231cce52724..be74306dacf4a3648e3227f11227e6399e6ed2eb 100644 --- a/deploy/hubserving/ocr_det/module.py +++ b/deploy/hubserving/ocr_det/module.py @@ -31,7 +31,7 @@ from tools.infer.predict_det import TextDetector author_email="paddle-dev@baidu.com", type="cv/text_recognition") class OCRDet(hub.Module): - def _initialize(self, use_gpu=False): + def _initialize(self, use_gpu=False, enable_mkldnn=False): """ initialize with the necessary elements """ @@ -51,6 +51,7 @@ class OCRDet(hub.Module): "Environment Variable CUDA_VISIBLE_DEVICES is not set correctly. If you wanna use gpu, please set CUDA_VISIBLE_DEVICES via export CUDA_VISIBLE_DEVICES=cuda_device_id." ) cfg.ir_optim = True + cfg.enable_mkldnn = enable_mkldnn self.text_detector = TextDetector(cfg) diff --git a/deploy/hubserving/ocr_rec/module.py b/deploy/hubserving/ocr_rec/module.py index cf612dd7376078f1badb41fe8655dccda9d3f05d..846f5437fe3b0a5136bff6c902481f888558d594 100644 --- a/deploy/hubserving/ocr_rec/module.py +++ b/deploy/hubserving/ocr_rec/module.py @@ -31,7 +31,7 @@ from tools.infer.predict_rec import TextRecognizer author_email="paddle-dev@baidu.com", type="cv/text_recognition") class OCRRec(hub.Module): - def _initialize(self, use_gpu=False): + def _initialize(self, use_gpu=False, enable_mkldnn=False): """ initialize with the necessary elements """ @@ -51,6 +51,7 @@ class OCRRec(hub.Module): "Environment Variable CUDA_VISIBLE_DEVICES is not set correctly. If you wanna use gpu, please set CUDA_VISIBLE_DEVICES via export CUDA_VISIBLE_DEVICES=cuda_device_id." ) cfg.ir_optim = True + cfg.enable_mkldnn = enable_mkldnn self.text_recognizer = TextRecognizer(cfg) diff --git a/deploy/hubserving/ocr_system/module.py b/deploy/hubserving/ocr_system/module.py index bed0c4e6fff372275aa415e60d878fe4dcd7acc4..cb526e1185d8eb623af84ba2451ff38523bbf642 100644 --- a/deploy/hubserving/ocr_system/module.py +++ b/deploy/hubserving/ocr_system/module.py @@ -31,7 +31,7 @@ from tools.infer.predict_system import TextSystem author_email="paddle-dev@baidu.com", type="cv/text_recognition") class OCRSystem(hub.Module): - def _initialize(self, use_gpu=False): + def _initialize(self, use_gpu=False, enable_mkldnn=False): """ initialize with the necessary elements """ @@ -51,7 +51,8 @@ class OCRSystem(hub.Module): "Environment Variable CUDA_VISIBLE_DEVICES is not set correctly. If you wanna use gpu, please set CUDA_VISIBLE_DEVICES via export CUDA_VISIBLE_DEVICES=cuda_device_id." ) cfg.ir_optim = True - + cfg.enable_mkldnn = enable_mkldnn + self.text_sys = TextSystem(cfg) def read_images(self, paths=[]):