From e33a952360a4eb54651d4aa4dba4f7fbac94dd45 Mon Sep 17 00:00:00 2001 From: MissPenguin Date: Wed, 15 Jul 2020 19:26:02 +0800 Subject: [PATCH] add parameter enable_mkldnn for hubserving modules --- deploy/hubserving/ocr_det/module.py | 3 ++- deploy/hubserving/ocr_rec/module.py | 3 ++- deploy/hubserving/ocr_system/module.py | 5 +++-- 3 files changed, 7 insertions(+), 4 deletions(-) diff --git a/deploy/hubserving/ocr_det/module.py b/deploy/hubserving/ocr_det/module.py index e5fac23f..be74306d 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 cf612dd7..846f5437 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 bed0c4e6..cb526e11 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=[]): -- GitLab