提交 3201bc6a 编写于 作者: S syyxsxx

add python mkl supported

上级 9cc2b119
......@@ -13,6 +13,7 @@
# limitations under the License.
import os
import os.path as osp
import psutil
import cv2
import numpy as np
import yaml
......@@ -30,7 +31,7 @@ class Predictor:
use_gpu=True,
gpu_id=0,
use_mkl=False,
mkl_thread_num=4,
mkl_thread_num=psutil.cpu_count(),
use_trt=False,
use_glog=False,
memory_optimize=True):
......@@ -84,7 +85,7 @@ class Predictor:
use_gpu=True,
gpu_id=0,
use_mkl=False,
mkl_thread_num=4,
mkl_thread_num=psutil.cpu_count(),
use_trt=False,
use_glog=False,
memory_optimize=True):
......@@ -98,6 +99,7 @@ class Predictor:
else:
config.disable_gpu()
if use_mkl:
if self.model_name not in ["HRNet", "DeepLabv3p"]:
config.enable_mkldnn()
config.set_cpu_math_library_num_threads(mkl_thread_num)
if use_glog:
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册