提交 e33a9523 编写于 作者: M MissPenguin

add parameter enable_mkldnn for hubserving modules

上级 17edc892
...@@ -31,7 +31,7 @@ from tools.infer.predict_det import TextDetector ...@@ -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): ...@@ -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)
......
...@@ -31,7 +31,7 @@ from tools.infer.predict_rec import TextRecognizer ...@@ -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): ...@@ -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)
......
...@@ -31,7 +31,7 @@ from tools.infer.predict_system import TextSystem ...@@ -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): ...@@ -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)
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册