diff --git a/deploy/cpp_infer/src/ocr_det.cpp b/deploy/cpp_infer/src/ocr_det.cpp index a935003bb453bcb13da2717eae642f1f09addf37..bf94abce236853410c15434d494058be03a62a81 100644 --- a/deploy/cpp_infer/src/ocr_det.cpp +++ b/deploy/cpp_infer/src/ocr_det.cpp @@ -26,6 +26,7 @@ void DBDetector::LoadModel(const std::string &model_dir) { config.DisableGpu(); if (this->use_mkldnn_) { config.EnableMKLDNN(); + // cache 10 different shapes for mkldnn to avoid memory leak config.SetMkldnnCacheCapacity(10); } config.SetCpuMathLibraryNumThreads(this->cpu_math_library_num_threads_); diff --git a/deploy/cpp_infer/src/ocr_rec.cpp b/deploy/cpp_infer/src/ocr_rec.cpp index 9d71d8c4553c280fb67022d2054680e93b767153..b997d8291a64f9b6042bce648bcd358e34d55a95 100644 --- a/deploy/cpp_infer/src/ocr_rec.cpp +++ b/deploy/cpp_infer/src/ocr_rec.cpp @@ -126,6 +126,7 @@ void CRNNRecognizer::LoadModel(const std::string &model_dir) { config.DisableGpu(); if (this->use_mkldnn_) { config.EnableMKLDNN(); + // cache 10 different shapes for mkldnn to avoid memory leak config.SetMkldnnCacheCapacity(10); } config.SetCpuMathLibraryNumThreads(this->cpu_math_library_num_threads_); diff --git a/deploy/cpp_infer/tools/config.txt b/deploy/cpp_infer/tools/config.txt index 40beea3a2e6f0260a42202d6411ffb10907bf871..6c53f29eeb310677815d106d3e0ae39fb03bc2e2 100644 --- a/deploy/cpp_infer/tools/config.txt +++ b/deploy/cpp_infer/tools/config.txt @@ -3,7 +3,7 @@ use_gpu 0 gpu_id 0 gpu_mem 4000 cpu_math_library_num_threads 10 -use_mkldnn 0 +use_mkldnn 1 use_zero_copy_run 1 # det config diff --git a/tools/infer/utility.py b/tools/infer/utility.py index 1d07f62f206ca2c250e340921e862bbccd5382b2..3e1f07b8a7127e64a994c34d296c945ad1cafd0a 100755 --- a/tools/infer/utility.py +++ b/tools/infer/utility.py @@ -101,6 +101,7 @@ def create_predictor(args, mode): config.disable_gpu() config.set_cpu_math_library_num_threads(6) if args.enable_mkldnn: + # cache 10 different shapes for mkldnn to avoid memory leak config.set_mkldnn_cache_capacity(10) config.enable_mkldnn()