提交 1076eda9 编写于 作者: J jiangjiajun

fix bug for GPU and MKLDNN

上级 d1935967
...@@ -66,7 +66,7 @@ void Model::create_predictor(const std::string& model_dir, ...@@ -66,7 +66,7 @@ void Model::create_predictor(const std::string& model_dir,
if (key == "") { if (key == "") {
config.SetModel(model_file, params_file); config.SetModel(model_file, params_file);
} }
if (use_mkl) { if (use_mkl && !use_gpu) {
if (name != "HRNet" && name != "DeepLabv3p" && name != "PPYOLO") { if (name != "HRNet" && name != "DeepLabv3p" && name != "PPYOLO") {
config.EnableMKLDNN(); config.EnableMKLDNN();
config.SetCpuMathLibraryNumThreads(mkl_thread_num); config.SetCpuMathLibraryNumThreads(mkl_thread_num);
......
...@@ -108,7 +108,7 @@ class Predictor: ...@@ -108,7 +108,7 @@ class Predictor:
config.enable_use_gpu(100, gpu_id) config.enable_use_gpu(100, gpu_id)
else: else:
config.disable_gpu() config.disable_gpu()
if use_mkl: if use_mkl and not use_gpu:
if self.model_name not in ["HRNet", "DeepLabv3p", "PPYOLO"]: if self.model_name not in ["HRNet", "DeepLabv3p", "PPYOLO"]:
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)
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册