瀏覽代碼

Fix memory leak when enabling mkldnn (#4148)

* Fix memory leak when enabling mkldnn

Fix https://github.com/PaddlePaddle/PaddleOCR/issues/15631

* Set default cache capacity to 10
timminator 5 月之前
父節點
當前提交
60ae2f544e
共有 2 個文件被更改,包括 10 次插入1 次删除
  1. 1 1
      paddlex/inference/models/common/static_infer.py
  2. 9 0
      paddlex/inference/utils/pp_option.py

+ 1 - 1
paddlex/inference/models/common/static_infer.py

@@ -411,7 +411,7 @@ class PaddleInfer(StaticInfer):
                         config.enable_mkldnn()
                         if "bf16" in self._option.run_mode:
                             config.enable_mkldnn_bfloat16()
-                        config.set_mkldnn_cache_capacity(-1)
+                        config.set_mkldnn_cache_capacity(self._option.mkldnn_cache_capacity)
                     else:
                         logging.warning(
                             "MKL-DNN is not available. We will disable MKL-DNN."

+ 9 - 0
paddlex/inference/utils/pp_option.py

@@ -119,6 +119,7 @@ class PaddlePredictorOption(object):
             "trt_dynamic_shape_input_data": None,  # only for trt
             "trt_shape_range_info_path": None,  # only for trt
             "trt_allow_rebuild_at_runtime": True,  # only for trt
+            "mkldnn_cache_capacity": 10,
         }
         return default_config
 
@@ -294,6 +295,14 @@ class PaddlePredictorOption(object):
     def trt_allow_rebuild_at_runtime(self, trt_allow_rebuild_at_runtime):
         self._update("trt_allow_rebuild_at_runtime", trt_allow_rebuild_at_runtime)
 
+    @property
+    def mkldnn_cache_capacity(self):
+        return self._cfg["mkldnn_cache_capacity"]
+
+    @mkldnn_cache_capacity.setter
+    def mkldnn_cache_capacity(self, capacity: int):
+        self._update("mkldnn_cache_capacity", capacity)
+
     # For backward compatibility
     # TODO: Issue deprecation warnings
     @property