提交 e33a9523 编写于 作者: M MissPenguin

add parameter enable_mkldnn for hubserving modules

上级 17edc892
......@@ -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)
......
......@@ -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)
......
......@@ -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,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."
)
cfg.ir_optim = True
cfg.enable_mkldnn = enable_mkldnn
self.text_sys = TextSystem(cfg)
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册