support paddle-trt

paddlex/inference/components/paddle_predictor/predictor.py

@@ -23,6 +23,29 @@ from ...utils.pp_option import PaddlePredictorOption
 from ..base import BaseComponent
 
 
+def collect_trt_shapes(
+    model_file, model_params, gpu_id, shape_range_info_path, trt_dynamic_shapes
+):
+    config = paddle.inference.Config(model_file, model_params)
+    config.enable_use_gpu(100, gpu_id)
+    min_arrs, opt_arrs, max_arrs = {}, {}, {}
+    for name, candidate_shapes in trt_dynamic_shapes.items():
+        min_shape, opt_shape, max_shape = candidate_shapes
+        min_arrs[name] = np.ones(min_shape, dtype=np.float32)
+        opt_arrs[name] = np.ones(opt_shape, dtype=np.float32)
+        max_arrs[name] = np.ones(max_shape, dtype=np.float32)
+
+    config.collect_shape_range_info(shape_range_info_path)
+    predictor = paddle.inference.create_predictor(config)
+    # opt_arrs would be used twice to simulate the most common situations
+    for arrs in [min_arrs, opt_arrs, opt_arrs, max_arrs]:
+        for name, arr in arrs.items():
+            input_handler = predictor.get_input_handle(name)
+            input_handler.reshape(arr.shape)
+            input_handler.copy_from_cpu(arr)
+        predictor.run()
+
+
 class Copy2GPU(BaseComponent):
 
     def __init__(self, input_handlers):
@@ -132,22 +155,25 @@ class BasePaddlePredictor(BaseComponent):
                         use_calib_mode=self.option.trt_calib_mode,
                     )
 
-                    if self.option.shape_info_filename is not None:
-                        if not os.path.exists(self.option.shape_info_filename):
-                            config.collect_shape_range_info(
-                                self.option.shape_info_filename
-                            )
-                            logging.info(
-                                f"Dynamic shape info is collected into: {self.option.shape_info_filename}"
-                            )
-                        else:
-                            logging.info(
-                                f"A dynamic shape info file ( {self.option.shape_info_filename} ) already exists. \
-        No need to generate again."
-                            )
-                        config.enable_tuned_tensorrt_dynamic_shape(
-                            self.option.shape_info_filename, True
+                    if not os.path.exists(self.option.shape_info_filename):
+                        logging.info(
+                            f"Dynamic shape info is collected into: {self.option.shape_info_filename}"
+                        )
+                        collect_trt_shapes(
+                            model_file,
+                            params_file,
+                            self.option.device_id,
+                            self.option.shape_info_filename,
+                            self.option.trt_dynamic_shapes,
                         )
+                    else:
+                        logging.info(
+                            f"A dynamic shape info file ( {self.option.shape_info_filename} ) already exists. No need to collect again."
+                        )
+                    config.enable_tuned_tensorrt_dynamic_shape(
+                        self.option.shape_info_filename, True
+                    )
+
         elif self.option.device == "npu":
             config.enable_custom_device("npu")
         elif self.option.device == "xpu":
@@ -196,9 +222,7 @@ class BasePaddlePredictor(BaseComponent):
 
         config.set_cpu_math_library_num_threads(self.option.cpu_threads)
 
-        if not (
-            self.option.device == "gpu" and self.option.run_mode.startswith("trt")
-        ):
+        if not (self.option.device == "gpu" and self.option.run_mode.startswith("trt")):
             if self.option.device in ("cpu", "gpu"):
                 if hasattr(config, "enable_new_ir"):
                     config.enable_new_ir(self.option.enable_new_ir)

paddlex/inference/models/base/basic_predictor.py

@@ -41,6 +41,14 @@ class BasicPredictor(
             pp_option = PaddlePredictorOption(model_name=self.model_name)
         if device:
             pp_option.device = device
+        trt_dynamic_shapes = (
+            self.config.get("Hpi", {})
+            .get("backend_configs", {})
+            .get("paddle_infer", {})
+            .get("trt_dynamic_shapes", None)
+        )
+        if trt_dynamic_shapes:
+            pp_option.trt_dynamic_shapes = trt_dynamic_shapes
         self.pp_option = pp_option
 
         self.components = {}

paddlex/inference/models_new/base/predictor/basic_predictor.py

@@ -54,6 +54,14 @@ class BasicPredictor(
             pp_option = PaddlePredictorOption(model_name=self.model_name)
         if device:
             pp_option.device = device
+        trt_dynamic_shapes = (
+            self.config.get("Hpi", {})
+            .get("backend_configs", {})
+            .get("paddle_infer", {})
+            .get("trt_dynamic_shapes", None)
+        )
+        if trt_dynamic_shapes:
+            pp_option.trt_dynamic_shapes = trt_dynamic_shapes
         self.pp_option = pp_option
 
         logging.debug(f"{self.__class__.__name__}: {self.model_dir}")

paddlex/inference/models_new/common/static_infer.py

@@ -24,6 +24,29 @@ from ....utils import logging
 from ...utils.pp_option import PaddlePredictorOption
 
 
+def collect_trt_shapes(
+    model_file, model_params, gpu_id, shape_range_info_path, trt_dynamic_shapes
+):
+    config = paddle.inference.Config(model_file, model_params)
+    config.enable_use_gpu(100, gpu_id)
+    min_arrs, opt_arrs, max_arrs = {}, {}, {}
+    for name, candidate_shapes in trt_dynamic_shapes.items():
+        min_shape, opt_shape, max_shape = candidate_shapes
+        min_arrs[name] = np.ones(min_shape, dtype=np.float32)
+        opt_arrs[name] = np.ones(opt_shape, dtype=np.float32)
+        max_arrs[name] = np.ones(max_shape, dtype=np.float32)
+
+    config.collect_shape_range_info(shape_range_info_path)
+    predictor = paddle.inference.create_predictor(config)
+    # opt_arrs would be used twice to simulate the most common situations
+    for arrs in [min_arrs, opt_arrs, opt_arrs, max_arrs]:
+        for name, arr in arrs.items():
+            input_handler = predictor.get_input_handle(name)
+            input_handler.reshape(arr.shape)
+            input_handler.copy_from_cpu(arr)
+        predictor.run()
+
+
 class Copy2GPU:
 
     def __init__(self, input_handlers):
@@ -147,22 +170,25 @@ class StaticInfer:
                         use_calib_mode=self.option.trt_calib_mode,
                     )
 
-                    if self.option.shape_info_filename is not None:
-                        if not os.path.exists(self.option.shape_info_filename):
-                            config.collect_shape_range_info(
-                                self.option.shape_info_filename
-                            )
-                            logging.info(
-                                f"Dynamic shape info is collected into: {self.option.shape_info_filename}"
-                            )
-                        else:
-                            logging.info(
-                                f"A dynamic shape info file ( {self.option.shape_info_filename} ) already exists. \
-        No need to generate again."
-                            )
-                        config.enable_tuned_tensorrt_dynamic_shape(
-                            self.option.shape_info_filename, True
+                    if not os.path.exists(self.option.shape_info_filename):
+                        logging.info(
+                            f"Dynamic shape info is collected into: {self.option.shape_info_filename}"
+                        )
+                        collect_trt_shapes(
+                            model_file,
+                            params_file,
+                            self.option.device_id,
+                            self.option.shape_info_filename,
+                            self.option.trt_dynamic_shapes,
                         )
+                    else:
+                        logging.info(
+                            f"A dynamic shape info file ( {self.option.shape_info_filename} ) already exists. No need to collect again."
+                        )
+                    config.enable_tuned_tensorrt_dynamic_shape(
+                        self.option.shape_info_filename, True
+                    )
+
         elif self.option.device == "npu":
             config.enable_custom_device("npu")
         elif self.option.device == "xpu":

paddlex/inference/utils/io/writers.py

@@ -357,7 +357,7 @@ class _BaseJsonWriterBackend(object):
 
 class JsonWriterBackend(_BaseJsonWriterBackend):
     def _write_obj(self, out_path, obj, **bk_args):
-        with open(out_path, "w") as f:
+        with open(out_path, "w", encoding="utf-8") as f:
             json.dump(obj, f, **bk_args)
 
 

paddlex/inference/utils/pp_option.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import os
+from typing import Dict, List
 
 from ...utils.device import parse_device, set_env_for_device, get_default_device
 from ...utils import logging
@@ -75,6 +76,7 @@ class PaddlePredictorOption(object):
             "delete_pass": [],
             "enable_new_ir": True if self.model_name not in NEWIR_BLOCKLIST else False,
             "batch_size": 1,  # only for trt
+            "trt_dynamic_shapes": {},  # only for trt
         }
 
     def _update(self, k, v):
@@ -150,6 +152,17 @@ class PaddlePredictorOption(object):
         self._update("shape_info_filename", shape_info_filename)
 
     @property
+    def trt_dynamic_shapes(self):
+        return self._cfg["trt_dynamic_shapes"]
+
+    @trt_dynamic_shapes.setter
+    def trt_dynamic_shapes(self, trt_dynamic_shapes: Dict[str, List[List[int]]]):
+        assert isinstance(trt_dynamic_shapes, dict)
+        for input_k in trt_dynamic_shapes:
+            assert isinstance(trt_dynamic_shapes[input_k], list)
+        self._update("trt_dynamic_shapes", trt_dynamic_shapes)
+
+    @property
     def trt_calib_mode(self):
         return self._cfg["trt_calib_mode"]