[Paddle-TRT]Clip op supports input type int (#66655)

ckl117 · web-flow · commit ef03571524bd · 2024-08-02T15:09:25.000+08:00
* trt clip op supports int32 input
diff --git a/paddle/fluid/inference/tensorrt/convert/clip_op.cc b/paddle/fluid/inference/tensorrt/convert/clip_op.cc
@@ -26,25 +26,36 @@ class ClipOpConverter : public OpConverter {
   void operator()(const framework::proto::OpDesc& op,
                   const framework::Scope& scope,
                   bool test_mode) override {
-#if IS_TRT_VERSION_GE(5130)
-    VLOG(3) << "convert a clip op to tensorrt IActivationLayer.";
+    VLOG(3) << "convert a clip op to tensorrt layer.";
     framework::OpDesc op_desc(op, nullptr);
     // Declare inputs
     auto* input = engine_->GetITensor(op_desc.Input("X")[0]);
     float min = PADDLE_GET_CONST(float, op_desc.GetAttr("min"));
     float max = PADDLE_GET_CONST(float, op_desc.GetAttr("max"));
-    auto* layer = TRT_ENGINE_ADD_LAYER(
-        engine_, Activation, *input, nvinfer1::ActivationType::kCLIP);
-    layer->setAlpha(min);
-    layer->setBeta(max);
+    int32_t rank = input->getDimensions().nbDims;
+    nvinfer1::ITensor* input_shape_tensor = Shape(input);
+    nvinfer1::DataType data_type = input->getType();
+    nvinfer1::ITensor* alphaT{nullptr};
+    nvinfer1::ITensor* betaT{nullptr};
+    if (data_type == nvinfer1::DataType::kINT32) {
+      alphaT =
+          FillConstantLayer(input_shape_tensor, rank, static_cast<int>(min));
+      betaT =
+          FillConstantLayer(input_shape_tensor, rank, static_cast<int>(max));
+    } else {
+      alphaT = FillConstantLayer(input_shape_tensor, rank, min);
+      betaT = FillConstantLayer(input_shape_tensor, rank, max);
+    }
+
+    auto* lowerClip = Max(input, alphaT);
+    auto* layer = TRT_ENGINE_ADD_LAYER(engine_,
+                                       ElementWise,
+                                       *lowerClip,
+                                       *betaT,
+                                       nvinfer1::ElementWiseOperation::kMIN);
 
     auto output_name = op_desc.Output("Out")[0];
     ReplenishLayerAndOutput(layer, "clip", {output_name}, test_mode);
-#else
-    PADDLE_THROW(
-        common::errors::Fatal("clip TRT converter is only supported on TRT "
-                              "5.1.3.0 or higher version."));
-#endif
   }
 };
 
diff --git a/paddle/fluid/inference/tensorrt/op_teller.cc b/paddle/fluid/inference/tensorrt/op_teller.cc
@@ -2254,6 +2254,11 @@ struct SimpleOpTypeSetTeller : public Teller {
     }
 
     if (op_type == "clip") {
+      if (!with_dynamic_shape) {
+        VLOG(3) << "the clip does not support static "
+                   "shape yet";
+        return false;
+      }
       // Paddle-TRT does not support the input tensors: Min and Max
       auto clip_inputs = desc.Inputs();
       if (clip_inputs.find("Min") != clip_inputs.end()) {
@@ -2277,15 +2282,9 @@ struct SimpleOpTypeSetTeller : public Teller {
       auto x_var_name = desc.Input("X")[0];
       auto* x_var_desc = block->FindVarRecursive(x_var_name);
       const auto x_shape = x_var_desc->GetShape();
-
-      auto dtype = x_var_desc->GetDataType();
-      if (dtype != framework::proto::VarType::FP32) {
-        return false;
-      }
-      if (!with_dynamic_shape && (x_shape.size() == 1 || x_shape.empty())) {
+      if (x_shape.empty()) {
         VLOG(3) << op_type
-                << " op does not support input's dim is 1 or 0 in tensorrt "
-                   "static shape mode.";
+                << " op does not support input's dim is 0 in tensorrt.";
         return false;
       }
     }
diff --git a/test/ir/inference/test_trt_convert_clip.py b/test/ir/inference/test_trt_convert_clip.py
@@ -28,17 +28,17 @@ def is_program_valid(self, program_config: ProgramConfig) -> bool:
         return True
 
     def sample_program_configs(self):
-        def generate_input1(dims, batch, attrs: List[Dict[str, Any]]):
+        def generate_input1(dims, batch, dtype, attrs: List[Dict[str, Any]]):
             if dims == 0:
-                return np.ones([]).astype(np.float32)
+                return np.ones([]).astype(dtype)
             elif dims == 1:
-                return np.ones([32]).astype(np.float32)
+                return np.ones([32]).astype(dtype)
             elif dims == 2:
-                return np.ones([3, 32]).astype(np.float32)
+                return np.ones([3, 32]).astype(dtype)
             elif dims == 3:
-                return np.ones([3, 32, 32]).astype(np.float32)
+                return np.ones([3, 32, 32]).astype(dtype)
             else:
-                return np.ones([batch, 3, 32, 32]).astype(np.float32)
+                return np.ones([batch, 3, 32, 32]).astype(dtype)
 
         def generate_weight1(attrs: List[Dict[str, Any]]):
             return np.array([np.random.uniform(1, 10)]).astype("float32")
@@ -48,50 +48,54 @@ def generate_weight2(attrs: List[Dict[str, Any]]):
 
         for dims in [0, 1, 2, 3, 4]:
             for batch in [1, 4]:
-                for op_inputs in [
-                    {"X": ["input_data"]},
-                    {"X": ["input_data"], "Min": ["Min_"], "Max": ["Max_"]},
-                ]:
-                    self.input_num = len(op_inputs)
-                    self.dims = dims
-                    dics = [
-                        {
-                            "min": np.random.uniform(1, 10),
-                            "max": np.random.uniform(10, 20),
-                        },
-                        {"op_inputs": op_inputs},
-                    ]
-                    ops_config = [
-                        {
-                            "op_type": "clip",
-                            "op_inputs": op_inputs,
-                            "op_outputs": {"Out": ["output_data"]},
-                            "op_attrs": dics[0],
-                        }
-                    ]
-                    ops = self.generate_op_config(ops_config)
-
-                    program_config = ProgramConfig(
-                        ops=ops,
-                        weights={
-                            "Min_": TensorConfig(
-                                data_gen=partial(generate_weight1, dics)
-                            ),
-                            "Max_": TensorConfig(
-                                data_gen=partial(generate_weight2, dics)
-                            ),
-                        },
-                        inputs={
-                            "input_data": TensorConfig(
-                                data_gen=partial(
-                                    generate_input1, dims, batch, dics
+                for dtype in [np.float32, np.int32]:
+                    for op_inputs in [
+                        {"X": ["input_data"]},
+                        {"X": ["input_data"], "Min": ["Min_"], "Max": ["Max_"]},
+                    ]:
+                        self.input_num = len(op_inputs)
+                        self.dims = dims
+                        dics = [
+                            {
+                                "min": np.random.uniform(1, 10),
+                                "max": np.random.uniform(10, 20),
+                            },
+                            {"op_inputs": op_inputs},
+                        ]
+                        ops_config = [
+                            {
+                                "op_type": "clip",
+                                "op_inputs": op_inputs,
+                                "op_outputs": {"Out": ["output_data"]},
+                                "op_attrs": dics[0],
+                            }
+                        ]
+                        ops = self.generate_op_config(ops_config)
+                        program_config = ProgramConfig(
+                            ops=ops,
+                            weights={
+                                "Min_": TensorConfig(
+                                    data_gen=partial(generate_weight1, dics)
+                                ),
+                                "Max_": TensorConfig(
+                                    data_gen=partial(generate_weight2, dics)
+                                ),
+                            },
+                            inputs={
+                                "input_data": TensorConfig(
+                                    data_gen=partial(
+                                        generate_input1,
+                                        dims,
+                                        batch,
+                                        dtype,
+                                        dics,
+                                    )
                                 )
-                            )
-                        },
-                        outputs=["output_data"],
-                    )
+                            },
+                            outputs=["output_data"],
+                        )
 
-                    yield program_config
+                        yield program_config
 
     def sample_predictor_configs(self, program_config):
         def generate_dynamic_shape(attrs):
@@ -128,13 +132,10 @@ def clear_dynamic_shape():
             self.dynamic_shape.opt_input_shape = {}
 
         def generate_trt_nodes_num(attrs, dynamic_shape):
-            if self.input_num == 3:
-                return 0, 3
+            if dynamic_shape and self.dims != 0 and self.input_num != 3:
+                return 1, 2
             else:
-                if not dynamic_shape and (self.dims == 1 or self.dims == 0):
-                    return 0, 3
-                else:
-                    return 1, 2
+                return 0, 3
 
         attrs = [
             program_config.ops[i].attrs for i in range(len(program_config.ops))
@@ -143,25 +144,21 @@ def generate_trt_nodes_num(attrs, dynamic_shape):
         # for static_shape
         clear_dynamic_shape()
         self.trt_param.precision = paddle_infer.PrecisionType.Float32
-        program_config.set_input_type(np.float32)
         yield self.create_inference_config(), generate_trt_nodes_num(
             attrs, False
         ), 1e-5
         self.trt_param.precision = paddle_infer.PrecisionType.Half
-        program_config.set_input_type(np.float16)
         yield self.create_inference_config(), generate_trt_nodes_num(
             attrs, False
         ), (1e-3, 1e-3)
 
         # for dynamic_shape
         generate_dynamic_shape(attrs)
         self.trt_param.precision = paddle_infer.PrecisionType.Float32
-        program_config.set_input_type(np.float32)
         yield self.create_inference_config(), generate_trt_nodes_num(
             attrs, True
         ), 1e-5
         self.trt_param.precision = paddle_infer.PrecisionType.Half
-        program_config.set_input_type(np.float16)
         yield self.create_inference_config(), generate_trt_nodes_num(
             attrs, True
         ), (1e-3, 1e-3)
