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

add python mkl supported

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