PaddlePaddle
diff --git a/‎python/paddle/tensorrt/__init__.py
Copy file name to clipboard
+2-2Lines changed: 2 additions & 2 deletions b/‎python/paddle/tensorrt/__init__.py
Copy file name to clipboard
+2-2Lines changed: 2 additions & 2 deletions
diff --git a/‎python/paddle/tensorrt/converter.py
Copy file name to clipboard
+76-48Lines changed: 76 additions & 48 deletions b/‎python/paddle/tensorrt/converter.py
Copy file name to clipboard
+76-48Lines changed: 76 additions & 48 deletions
diff --git a/‎python/paddle/tensorrt/converter_utils.py
Copy file name to clipboard
+6Lines changed: 6 additions & 0 deletions b/‎python/paddle/tensorrt/converter_utils.py
Copy file name to clipboard
+6Lines changed: 6 additions & 0 deletions
@@ -12,11 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .export import Input, TensorRTConfig, convert, convert_loaded_model
+from .export import Input, PrecisionMode, TensorRTConfig, convert
 
 __all__ = [
     'Input',
     'TensorRTConfig',
     'convert',
-    'convert_loaded_model',
+    'PrecisionMode',
 ]
@@ -15,7 +15,6 @@
 import ctypes
 import hashlib
 import logging
-import os
 
 import numpy as np
 import tensorrt as trt
@@ -48,65 +47,29 @@
 from .impls.vision import *  # noqa: F403
 from .register import converter_registry
 from .util import (
+    TensorRTConfigManager,
+    get_cache_path,
+    get_trt_version,
     get_trt_version_list,
+    is_shape_tensor,
     map_dtype,
+    remove_duplicate_value,
     weight_to_tensor,
     zero_dims_to_one_dims,
 )
 
 version_list = get_trt_version_list()
 
-
-def get_cache_path():
-    home_path = os.path.expanduser("~")
-    cache_path = os.path.join(home_path, ".pp_trt_cache")
-
-    if not os.path.exists(cache_path):
-        os.makedirs(cache_path)
-    return cache_path
-
-
 _logger = get_logger(
     __name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s'
 )
 
 
-def get_trt_version():
-    return trt.__version__
-
-
-def remove_duplicate_value(value_list):
-    ret_list = []
-    ret_list_id = []
-    for value in value_list:
-        if value.id not in ret_list_id:
-            ret_list.append(value)
-            ret_list_id.append(value.id)
-    return ret_list
-
-
-# We use a special rule to judge whether a paddle value is a shape tensor.
-# The rule is consistent with the rule in C++ source code(collect_shape_manager.cc).
-# We use the rule for getting min/max/opt value shape from collect_shape_manager.
-# We don't use trt_tensor.is_shape_tensor, because sometimes, the trt_tensor that corresponding to paddle value is not a shape tensor
-# when it is a output in this trt graph, but it is a shape tensor when it is a input in next trt graph.
-def is_shape_tensor(value):
-    dims = value.shape
-    total_elements = 1
-    if (
-        dims.count(-1) > 1
-    ):  # we can only deal with the situation that is has one dynamic dims
-        return False
-    for dim in dims:
-        total_elements *= abs(dim)  # add abs for dynamic shape -1
-    is_int_dtype = value.dtype == paddle.int32 or value.dtype == paddle.int64
-    return total_elements <= 8 and total_elements >= 1 and is_int_dtype
-
-
 class PaddleToTensorRTConverter:
-    def __init__(self, paddle_program, scope):
+    def __init__(self, paddle_program, scope, trt_config=None):
         self.scope = scope
         self.program = paddle_program
+        self.trt_config = trt_config
         params = paddle_program.global_block().all_parameters()
         param_dict = {}
         # save parameters
@@ -116,6 +79,17 @@ def __init__(self, paddle_program, scope):
             # weights = trt.Weights(weight_array)
             param_dict.update({name: weight_array})
         self.param_dict = param_dict
+
+        trt_manager = TensorRTConfigManager()
+        if (
+            self.trt_config is not None
+            and self.trt_config.tensorrt_ops_run_float
+        ):
+            trt_manager.set_force_fp32_ops(
+                self.trt_config.tensorrt_ops_run_float
+            )
+            _logger.info(f"force_fp32_ops: {trt_manager.get_force_fp32_ops()}")
+
         self.input_info = {}
         self.trt_output_value_map = {}
 
@@ -157,6 +131,8 @@ def __is_output_value(value):
         return input_values, graph_output_values
 
     def convert_subgraph_to_trt(self, program, group_op):
+        from .export import PrecisionMode
+
         _logger.info(f"start process {group_op}")
 
         operations = next(iter(group_op.blocks())).ops
@@ -282,7 +258,7 @@ def convert_subgraph_to_trt(self, program, group_op):
                     value_to_trt_tensor[result.id] = None
 
         # Set TRT min/opt/max input shape and the value of shape tensor
-        for value in origin_input_value:
+        for i, value in enumerate(origin_input_value):
             trt_input = value_to_trt_tensor[value.id]
             if isinstance(trt_input, trt.Weights):
                 continue
@@ -324,6 +300,7 @@ def convert_subgraph_to_trt(self, program, group_op):
                     max_shape = get_value_shape_range_info(
                         value, False, paddle.base.core.ShapeMode.kMAX
                     )
+
                     if trt_input.is_shape_tensor:
                         min_value = get_value_shape_range_info(
                             value, True, paddle.base.core.ShapeMode.kMIN
@@ -419,7 +396,58 @@ def convert_subgraph_to_trt(self, program, group_op):
         ):  # trt version >= 8.6
             config.builder_optimization_level = 5
         config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)
-        trt_engine = builder.build_engine(network, config)
+
+        if self.trt_config is not None:
+            precision_mode = self.trt_config.precision_mode
+        if self.trt_config is not None and precision_mode == PrecisionMode.FP16:
+            if builder.platform_has_fast_fp16:
+                config.set_flag(trt.BuilderFlag.FP16)
+                _logger.info("Run Paddle-TRT FP16 mode")
+            else:
+                _logger.warning(
+                    "Hardware does not support FP16. Continuing in FP32 mode."
+                )
+        elif (
+            self.trt_config is not None and precision_mode == PrecisionMode.BF16
+        ):
+            if version_list[0] >= 9:
+                if builder.platform_has_fast_bfp16 and hasattr(
+                    builder, 'plateform_has_fast_bf16'
+                ):
+                    config.set_flag(trt.BuilderFlag.BF16)
+                    _logger.info("Run Paddle-TRT BF16 mode")
+                else:
+                    _logger.warning(
+                        "Hardware does not support BF16. Continuing in FP32 mode."
+                    )
+            else:
+                if builder.platform_has_fast_fp16:
+                    config.set_flag(trt.BuilderFlag.FP16)
+                    _logger.warning(
+                        "Because the version of TensorRT is less than 9.0, run  Paddle-TRT FP16 mode"
+                    )
+                else:
+                    _logger.warning(
+                        "Hardware does not support FP16. Continuing in FP32 mode."
+                    )
+        elif self.trt_config is not None:
+            _logger.info(
+                f"Default precision mode {self.trt_config.precision_mode}"
+            )
+
+        if (
+            version_list[0] > 8
+            or version_list[0] == 8
+            and version_list[1] >= 2
+            and version_list[2] >= 1
+        ):
+            if (
+                self.trt_config is not None
+                and self.trt_config.tensorrt_ops_run_float
+            ):
+                config.set_flag(trt.BuilderFlag.PREFER_PRECISION_CONSTRAINTS)
+
+        trt_engine = builder.build_serialized_network(network, config)
         trt_params = paddle.base.libpaddle.TRTEngineParams()
         trt_params.min_input_shape = min_shape_map
         trt_params.max_input_shape = max_shape_map
@@ -435,7 +463,7 @@ def convert_subgraph_to_trt(self, program, group_op):
         CACHE_ROOT = get_cache_path()
         CACHE_FILE = f"{CACHE_ROOT}/engine_{engine_name}.trt"
         with open(CACHE_FILE, "wb") as f:
-            f.write(trt_engine.serialize())
+            f.write(trt_engine)
         PIR_DUMP_FILE = f"{CACHE_ROOT}/engine_{engine_name}.pir"
         with open(PIR_DUMP_FILE, "w") as f:
             f.write(group_str)
@@ -503,5 +531,5 @@ def convert_program_to_trt(self):
                     orin_out_values[o_i].replace_all_uses_with(new_out[o_i])
 
                 self.program.global_block().remove_op(op)
-        # # Call clear_shape_info to clear the previous shape information
+        # Call clear_shape_info to clear the previous shape information
         clear_shape_info()
@@ -136,6 +136,8 @@ def get_positive_dim(dim, dim_size):
 
 
 def add_elementwise_layer(network, paddle_op, inputs, op_type):
+    from paddle.tensorrt.util import support_fp32_mix_precision
+
     weight_shape = paddle_op.operands()[1].source().shape
     input_shape = paddle_op.operands()[0].source().shape
 
@@ -157,6 +159,7 @@ def add_elementwise_layer(network, paddle_op, inputs, op_type):
         weight_tensor.name,
     )
     layer = network.add_elementwise(lhs_val, rhs_val, op_type)
+    support_fp32_mix_precision(paddle_op.name(), layer)
     return layer.get_output(0)
 
 
@@ -414,6 +417,8 @@ def trt_reduce_to_scalar(network, tensor):
 
 
 def convert_conv2d(network, paddle_op, inputs):
+    from paddle.tensorrt.util import support_fp32_mix_precision
+
     if (
         paddle_op.name() == "pd_op.conv2d"
         or paddle_op.name() == "pd_op.depthwise_conv2d"
@@ -520,6 +525,7 @@ def convert_conv2d(network, paddle_op, inputs):
         nv_dilations = trt.DimsHW(1, 1)
 
     layer.dilation_nd = nv_dilations
+    support_fp32_mix_precision(paddle_op.name(), layer)
 
     return layer.get_output(0)