add inference api:exp_disable_tensorrt_subgraph (#61967)

lizexu123 · web-flow · commit 365d8be7512e · 2024-02-27T16:21:11.000+08:00
diff --git a/paddle/fluid/framework/ir/subgraph_detector.cc b/paddle/fluid/framework/ir/subgraph_detector.cc
@@ -424,6 +424,20 @@ void SubGraphFuser::ReplaceNodesWithSubGraphs() {
   auto subgraphs = SubgraphDetector(graph_, node_inside_subgraph_teller_)();
   for (auto &subgraph : subgraphs) {
     if (subgraph.size() <= static_cast<size_t>(min_subgraph_size_)) continue;
+
+    bool continue_run = true;
+
+    for (auto *node : subgraph) {
+      for (const auto tmp_name : node->outputs) {
+        if (std::find(trt_exclude_var_names_.begin(),
+                      trt_exclude_var_names_.end(),
+                      tmp_name->Name()) != trt_exclude_var_names_.end()) {
+          continue_run = false;
+        }
+      }
+    }
+
+    if (continue_run == false) continue;
     std::unordered_set<Node *> subgraph_uniq(subgraph.begin(), subgraph.end());
     // replace this sub-graph with the first node. Two steps: 1. Create a Block
     // Node that contains this subgraph 2. Mark the nodes inside the sub-graph
diff --git a/paddle/fluid/framework/ir/subgraph_detector.h b/paddle/fluid/framework/ir/subgraph_detector.h
@@ -71,10 +71,12 @@ class SubGraphFuser {
   SubGraphFuser(Graph *graph,
                 const NodeInsideSubgraphTeller &teller,
                 int min_subgraph_size,
+                const std::vector<std::string> &trt_exclude_var_names = {},
                 std::string name = "tensorrt_engine")
       : graph_(graph),
         node_inside_subgraph_teller_(teller),
         min_subgraph_size_{min_subgraph_size},
+        trt_exclude_var_names_(trt_exclude_var_names),
         name_{name} {}
 
   // The main method which run all the logic.
@@ -88,6 +90,7 @@ class SubGraphFuser {
   Graph *graph_;
   NodeInsideSubgraphTeller node_inside_subgraph_teller_;
   int min_subgraph_size_;
+  std::vector<std::string> trt_exclude_var_names_;
   const std::string name_;
 };
 
diff --git a/paddle/fluid/inference/analysis/argument.h b/paddle/fluid/inference/analysis/argument.h
@@ -247,6 +247,9 @@ struct Argument {
   DECL_ARGUMENT_FIELD(trt_output_tensor_names,
                       TRTOutputTensorNames,
                       std::vector<std::string>);
+  DECL_ARGUMENT_FIELD(trt_exclude_var_names,
+                      TRTExcludeVarNames,
+                      std::vector<std::string>);
   DECL_ARGUMENT_FIELD(tensorrt_disabled_ops,
                       TensorRtDisabledOPs,
                       std::vector<std::string>);
diff --git a/paddle/fluid/inference/analysis/ir_pass_manager.cc b/paddle/fluid/inference/analysis/ir_pass_manager.cc
@@ -170,6 +170,9 @@ void IRPassManager::CreatePasses(Argument *argument,
       pass->Set(
           "output_tensor_names",
           new std::vector<std::string>(argument->trt_output_tensor_names()));
+      pass->Set(
+          "trt_exclude_var_names",
+          new std::vector<std::string>(argument->trt_exclude_var_names()));
       pass->Set("program",
                 new framework::ProgramDesc *(&argument->main_program()));
       pass->Set("predictor_id", new int(argument->predictor_id()));
diff --git a/paddle/fluid/inference/analysis/ir_passes/dlnne_subgraph_pass.cc b/paddle/fluid/inference/analysis/ir_passes/dlnne_subgraph_pass.cc
@@ -349,6 +349,7 @@ void analysis::DlnneSubgraphPass::ApplyImpl(framework::ir::Graph *graph) const {
       graph,
       teller,
       Get<int>("min_subgraph_size") /*min subgraph size*/,
+      {},
       "dlnne_engine");
   fuser();
 
diff --git a/paddle/fluid/inference/analysis/ir_passes/lite_subgraph_pass.cc b/paddle/fluid/inference/analysis/ir_passes/lite_subgraph_pass.cc
@@ -429,7 +429,8 @@ void LiteSubgraphPass::ApplyImpl(framework::ir::Graph* graph) const {
                                                     *node->Op());
   };
 
-  SubGraphFuser fuser(graph, teller, 0 /* min_subgraph_size */, "lite_engine");
+  SubGraphFuser fuser(
+      graph, teller, 0 /* min_subgraph_size */, {}, "lite_engine");
   fuser();
 
   std::vector<std::string> repetitive_params;
diff --git a/paddle/fluid/inference/analysis/ir_passes/tensorrt_subgraph_pass.cc b/paddle/fluid/inference/analysis/ir_passes/tensorrt_subgraph_pass.cc
@@ -183,6 +183,7 @@ void analysis::TensorRtSubgraphPass::ApplyImpl(
       graph,
       teller,
       Get<int>("min_subgraph_size") /*min subgraph size*/,
+      Get<std::vector<std::string>>("trt_exclude_var_names"),
       "tensorrt_engine");
   fuser();
 
diff --git a/paddle/fluid/inference/api/analysis_config.cc b/paddle/fluid/inference/api/analysis_config.cc
@@ -484,6 +484,7 @@ AnalysisConfig::AnalysisConfig(const AnalysisConfig &other) {
   CP_MEMBER(trt_engine_memory_sharing_identifier_);
   CP_MEMBER(trt_optimization_level_);
   CP_MEMBER(trt_ops_run_float_);
+  CP_MEMBER(trt_exclude_var_names_);
   // Dlnne related
   CP_MEMBER(use_dlnne_);
   CP_MEMBER(dlnne_min_subgraph_size_);
@@ -865,6 +866,13 @@ void AnalysisConfig::Exp_DisableTensorRtOPs(
   trt_disabled_ops_.insert(trt_disabled_ops_.end(), ops.begin(), ops.end());
 }
 
+void AnalysisConfig::Exp_DisableTensorRtSubgraph(
+    const std::vector<std::string> &var_name_not_trt) {
+  trt_exclude_var_names_.insert(trt_exclude_var_names_.end(),
+                                var_name_not_trt.begin(),
+                                var_name_not_trt.end());
+}
+
 void AnalysisConfig::EnableVarseqlen() { trt_use_varseqlen_ = true; }
 
 void AnalysisConfig::SetTensorRtOptimizationLevel(int level) {
@@ -1127,6 +1135,9 @@ std::string AnalysisConfig::SerializeInfoCache() {
   for (auto &op : trt_disabled_ops_) ss << op.c_str();
   ss << ";";
 
+  for (auto &name : trt_exclude_var_names_) ss << name.c_str();
+  ss << ";";
+
   ss << trt_use_dla_;
   ss << trt_dla_core_;
 
diff --git a/paddle/fluid/inference/api/analysis_predictor.cc b/paddle/fluid/inference/api/analysis_predictor.cc
@@ -1702,6 +1702,7 @@ void AnalysisPredictor::PrepareArgument() {
     argument_->SetTRTMarkOutput(config_.trt_mark_output_);
     argument_->SetTRTOutputTensorNames(config_.trt_output_tensor_names_);
     argument_->SetTensorRtDisabledOPs(config_.trt_disabled_ops_);
+    argument_->SetTRTExcludeVarNames(config_.trt_exclude_var_names_);
     argument_->SetTensorRtUseDLA(config_.trt_use_dla_);
     argument_->SetTensorRtDLACore(config_.trt_dla_core_);
     argument_->SetTensorRtUseStaticEngine(config_.trt_use_static_engine_);
diff --git a/paddle/fluid/inference/api/paddle_analysis_config.h b/paddle/fluid/inference/api/paddle_analysis_config.h
@@ -810,6 +810,9 @@ struct PD_INFER_DECL AnalysisConfig {
   ///
   void Exp_DisableTensorRtOPs(const std::vector<std::string>& ops);
 
+  void Exp_DisableTensorRtSubgraph(
+      const std::vector<std::string>& var_name_not_trt);
+
   ///
   /// \brief Replace some TensorRT plugins to TensorRT OSS(
   /// https://github.com/NVIDIA/TensorRT), with which some models's inference
@@ -1269,6 +1272,7 @@ struct PD_INFER_DECL AnalysisConfig {
   bool trt_with_interleaved_{false};
   bool trt_mark_output_{false};
   std::vector<std::string> trt_output_tensor_names_{};
+  std::vector<std::string> trt_exclude_var_names_{};
   std::string tensorrt_transformer_posid_{""};
   std::string tensorrt_transformer_maskid_{""};
   bool trt_use_dla_{false};
diff --git a/paddle/fluid/pybind/inference_api.cc b/paddle/fluid/pybind/inference_api.cc
@@ -20,7 +20,6 @@
 
 #include <cstring>
 #include <functional>
-#include <iostream>
 #include <iterator>
 #include <map>
 #include <memory>
@@ -933,6 +932,8 @@ void BindAnalysisConfig(py::module *m) {
       .def("trt_allow_build_at_runtime",
            &AnalysisConfig::trt_allow_build_at_runtime)
       .def("exp_disable_tensorrt_ops", &AnalysisConfig::Exp_DisableTensorRtOPs)
+      .def("exp_disable_tensorrt_subgraph",
+           &AnalysisConfig::Exp_DisableTensorRtSubgraph)
       .def("enable_tensorrt_dla",
            &AnalysisConfig::EnableTensorRtDLA,
            py::arg("dla_core") = 0)
diff --git a/test/ir/inference/test_trt_exp_tensorrt_subgraph.py b/test/ir/inference/test_trt_exp_tensorrt_subgraph.py
@@ -0,0 +1,190 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+from functools import partial
+
+import numpy as np
+from program_config import ProgramConfig, TensorConfig
+from trt_layer_auto_scan_test import TrtLayerAutoScanTest
+
+import paddle.inference as paddle_infer
+
+
+class TrtConvertSetValue(TrtLayerAutoScanTest):
+    def is_program_valid(self, program_config: ProgramConfig) -> bool:
+        return True
+
+    def create_inference_config(self, use_trt=True) -> paddle_infer.Config:
+        config = paddle_infer.Config()
+        config.disable_glog_info()
+        config.enable_use_gpu(100, 0)
+        config.exp_disable_tensorrt_subgraph(["input_data"])
+        config.set_optim_cache_dir(self.cache_dir)
+        if use_trt:
+            config.switch_ir_debug()
+            config.enable_tensorrt_engine(
+                max_batch_size=self.trt_param.max_batch_size,
+                workspace_size=self.trt_param.workspace_size,
+                min_subgraph_size=self.trt_param.min_subgraph_size,
+                precision_mode=self.trt_param.precision,
+                use_static=self.trt_param.use_static,
+                use_calib_mode=self.trt_param.use_calib_mode,
+            )
+            if self.dynamic_shape.min_input_shape and (
+                self.dynamic_shape.min_input_shape.keys()
+                == self.dynamic_shape.max_input_shape.keys()
+                == self.dynamic_shape.opt_input_shape.keys()
+            ):
+                config.set_trt_dynamic_shape_info(
+                    self.dynamic_shape.min_input_shape,
+                    self.dynamic_shape.max_input_shape,
+                    self.dynamic_shape.opt_input_shape,
+                    self.dynamic_shape.disable_trt_plugin_fp16,
+                )
+        return config
+
+    def sample_program_configs(self):
+        def generate_input1():
+            return np.random.random([2, 3, 3]).astype(np.float32)
+
+        def generate_input2():
+            return np.random.random([2, 2, 3]).astype(np.float32)
+
+        for update_scalar in [True, False]:
+            self.update_scalar = update_scalar
+            set_value_inputs = {}
+            if update_scalar:
+                set_value_inputs = {
+                    "Input": ["input_data"],
+                }
+            else:
+                set_value_inputs = {
+                    "Input": ["input_data"],
+                    "ValueTensor": ["update_data"],
+                }
+            ops_config = [
+                {
+                    "op_type": "set_value",
+                    "op_inputs": set_value_inputs,
+                    "op_outputs": {"Out": ["input_data"]},
+                    "op_attrs": {
+                        "axes": [1],
+                        "starts": [0],
+                        "ends": [2],
+                        "steps": [1],
+                        "decrease_axes": [],
+                        "values": [0.0],
+                    },
+                },
+                {
+                    "op_type": "relu",
+                    "op_inputs": {
+                        "X": ["input_data"],
+                    },
+                    "op_outputs": {"Out": ["output_data"]},
+                    "op_attrs": {},
+                },
+            ]
+
+            ops = self.generate_op_config(ops_config)
+            if update_scalar:
+                program_config = ProgramConfig(
+                    ops=ops,
+                    weights={},
+                    inputs={
+                        "input_data": TensorConfig(
+                            data_gen=partial(generate_input1)
+                        ),
+                    },
+                    outputs=["output_data"],
+                )
+            else:
+                program_config = ProgramConfig(
+                    ops=ops,
+                    weights={},
+                    inputs={
+                        "input_data": TensorConfig(
+                            data_gen=partial(generate_input1)
+                        ),
+                        "update_data": TensorConfig(
+                            data_gen=partial(generate_input2)
+                        ),
+                    },
+                    outputs=["output_data"],
+                )
+
+            yield program_config
+
+    def sample_predictor_configs(self, program_config):
+        def generate_dynamic_shape(attrs):
+            if self.update_scalar:
+                self.dynamic_shape.min_input_shape = {
+                    "input_data": [2, 3, 3],
+                }
+                self.dynamic_shape.max_input_shape = {
+                    "input_data": [3, 3, 4],
+                }
+                self.dynamic_shape.opt_input_shape = {
+                    "input_data": [3, 3, 3],
+                }
+            else:
+                self.dynamic_shape.min_input_shape = {
+                    "input_data": [2, 3, 3],
+                    "update_data": [2, 2, 3],
+                }
+                self.dynamic_shape.max_input_shape = {
+                    "input_data": [3, 3, 4],
+                    "update_data": [3, 2, 4],
+                }
+                self.dynamic_shape.opt_input_shape = {
+                    "input_data": [3, 3, 3],
+                    "update_data": [3, 2, 3],
+                }
+
+        def clear_dynamic_shape():
+            self.dynamic_shape.max_input_shape = {}
+            self.dynamic_shape.min_input_shape = {}
+            self.dynamic_shape.opt_input_shape = {}
+
+        def generate_trt_nodes_num(attrs, dynamic_shape):
+            if dynamic_shape:
+                ver = paddle_infer.get_trt_compile_version()
+                if self.update_scalar:
+                    if ver[0] * 1000 + ver[1] * 100 + ver[2] * 10 < 8200:
+                        return 1, 3
+                    return 0, 4
+                else:
+                    if ver[0] * 1000 + ver[1] * 100 + ver[2] * 10 < 8200:
+                        return 1, 4
+                    return 0, 5
+
+        attrs = [
+            program_config.ops[i].attrs for i in range(len(program_config.ops))
+        ]
+
+        generate_dynamic_shape(attrs)
+        self.trt_param.precision = paddle_infer.PrecisionType.Float32
+        # program_config.set_input_type(np.float32)
+        self.trt_param.workspace_size = 2013265920
+        yield self.create_inference_config(), generate_trt_nodes_num(
+            attrs, True
+        ), (1e-5, 1e-4)
+
+    def test(self):
+        self.run_test()
+
+
+if __name__ == "__main__":
+    unittest.main()
