forgetest5
/
PaddleOCR
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

add parameter enable_mkldnn for hubserving modules

This commit is contained in:
MissPenguin 2020-07-15 19:26:02 +08:00
parent 17edc892c5
commit e33a952360
3 changed files with 7 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
deploy/hubserving/ocr_det/module.py
View File

@ -31,7 +31,7 @@ from tools.infer.predict_det import TextDetector
author_email="paddle-dev@baidu.com", author_email="paddle-dev@baidu.com",
type="cv/text_recognition") type="cv/text_recognition")
class OCRDet(hub.Module): class OCRDet(hub.Module):
def _initialize(self, use_gpu=False): def _initialize(self, use_gpu=False, enable_mkldnn=False):
""" """
initialize with the necessary elements 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." "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.ir_optim = True
cfg.enable_mkldnn = enable_mkldnn
self.text_detector = TextDetector(cfg) self.text_detector = TextDetector(cfg)

3
deploy/hubserving/ocr_rec/module.py
View File

@ -31,7 +31,7 @@ from tools.infer.predict_rec import TextRecognizer
author_email="paddle-dev@baidu.com", author_email="paddle-dev@baidu.com",
type="cv/text_recognition") type="cv/text_recognition")
class OCRRec(hub.Module): class OCRRec(hub.Module):
def _initialize(self, use_gpu=False): def _initialize(self, use_gpu=False, enable_mkldnn=False):
""" """
initialize with the necessary elements 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." "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.ir_optim = True
cfg.enable_mkldnn = enable_mkldnn
self.text_recognizer = TextRecognizer(cfg) self.text_recognizer = TextRecognizer(cfg)

3
deploy/hubserving/ocr_system/module.py
View File

@ -31,7 +31,7 @@ from tools.infer.predict_system import TextSystem
author_email="paddle-dev@baidu.com", author_email="paddle-dev@baidu.com",
type="cv/text_recognition") type="cv/text_recognition")
class OCRSystem(hub.Module): class OCRSystem(hub.Module):
def _initialize(self, use_gpu=False): def _initialize(self, use_gpu=False, enable_mkldnn=False):
""" """
initialize with the necessary elements initialize with the necessary elements
""" """
@ -51,6 +51,7 @@ 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." "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.ir_optim = True
cfg.enable_mkldnn = enable_mkldnn
self.text_sys = TextSystem(cfg) self.text_sys = TextSystem(cfg)