Add support for conv2d explicit quantization (#57160)

leo0519 · web-flow · commit 012eca1866cc · 2023-09-27T14:18:12.000+08:00
* Add support for conv2d explicit quantization. Add interface to check whether use explicit quantization or not in TRT engine
diff --git a/paddle/fluid/framework/ir/constant_folding_pass.cc b/paddle/fluid/framework/ir/constant_folding_pass.cc
@@ -64,7 +64,11 @@ void ConstantFoldingPass::ApplyImpl(ir::Graph *graph) const {
       platform::errors::Fatal(
           "scope must not be null when applying constant folding."));
 
-  std::vector<std::string> blacklist{"feed", "matrix_multiply", "save"};
+  std::vector<std::string> blacklist{"feed",
+                                     "matrix_multiply",
+                                     "save",
+                                     "quantize_linear",
+                                     "dequantize_linear"};
   int folded_op_num = 0;
 
   auto op_node_sorted = framework::ir::TopologyVarientSort(
diff --git a/paddle/fluid/inference/analysis/ir_pass_manager.cc b/paddle/fluid/inference/analysis/ir_pass_manager.cc
@@ -226,6 +226,8 @@ void IRPassManager::CreatePasses(Argument *argument,
       pass->Set("use_inspector", new bool(argument->tensorrt_use_inspector()));
       pass->Set("inspector_serialize",
                 new bool(argument->tensorrt_inspector_serialize()));
+      pass->Set("use_explicit_quantization",
+                new bool(argument->tensorrt_use_explicit_quantization()));
 
       // tuned trt dynamic_shape
       pass->Set("trt_shape_range_info_path",
diff --git a/paddle/fluid/inference/analysis/ir_passes/tensorrt_subgraph_pass.cc b/paddle/fluid/inference/analysis/ir_passes/tensorrt_subgraph_pass.cc
@@ -143,6 +143,7 @@ void analysis::TensorRtSubgraphPass::ApplyImpl(
   bool no_calib_int8 = enable_int8 && !(use_calib_mode);
   auto trt_disabled_ops = Get<std::vector<std::string>>("trt_disabled_ops");
   auto with_dynamic_shape = Get<bool>("with_dynamic_shape");
+  auto use_explicit_quantization = Get<bool>("use_explicit_quantization");
   auto teller = [&](const framework::ir::Node *node) {
     if (!node->IsOp() || !node->Op()) return false;
     if (find(trt_disabled_ops.begin(),
@@ -163,7 +164,7 @@ void analysis::TensorRtSubgraphPass::ApplyImpl(
       }
     }
     bool is_ok = tensorrt::OpTeller::Global().Tell(
-        node, no_calib_int8, with_dynamic_shape);
+        node, no_calib_int8, with_dynamic_shape, use_explicit_quantization);
     if (!is_ok)
       VLOG(3) << node->Op()->Type().c_str() << " op is not in TensorRT";
     return is_ok;
@@ -572,6 +573,7 @@ std::string TensorRtSubgraphPass::CreateTensorRTOp(
   auto workspace_size = Get<int64_t>("workspace_size");
   auto gpu_device_id = Get<int>("gpu_device_id");
   auto optimization_level = Get<int>("optimization_level");
+  auto use_explicit_quantization = Get<bool>("use_explicit_quantization");
 
   // Set op's attrs.
   op_desc->SetType("tensorrt_engine");
@@ -775,6 +777,7 @@ std::string TensorRtSubgraphPass::CreateTensorRTOp(
   params.engine_info_path = engine_info_path;
   params.enable_low_precision_io = enable_low_precision_io;
   params.optimization_level = optimization_level;
+  params.use_explicit_quantization = use_explicit_quantization;
 
   tensorrt::TensorRTEngine *trt_engine =
       inference::Singleton<inference::tensorrt::TRTEngineManager>::Global()
diff --git a/paddle/fluid/inference/tensorrt/convert/conv2d_op.cc b/paddle/fluid/inference/tensorrt/convert/conv2d_op.cc
@@ -43,16 +43,45 @@ void ConvertConv2d(TensorRTEngine* engine,
   framework::OpDesc op_desc(op, nullptr);
 
   auto* X = engine->GetITensor(op_desc.Input("Input").front());
-  std::string filter_var_name = op_desc.Input("Filter").front();
-  auto* Y_v = scope.FindVar(filter_var_name);
-  phi::DenseTensor* Y_t = nullptr;
-  nvinfer1::ITensor* filter = nullptr;
+  bool enable_int8 = op_desc.HasAttr("enable_int8");
+
+  if (enable_int8) {
+#if IS_TRT_VERSION_GE(5000)
+    float in_scale = PADDLE_GET_CONST(float, op_desc.GetAttr("Input_scale"));
+    engine->SetTensorDynamicRange(X, in_scale);
+#endif
+  }
+
   int n_output;
   int n_input;
   int filter_h;
   int filter_w;
-  if (Y_v) {
-    Y_t = Y_v->GetMutable<phi::DenseTensor>();
+  std::string filter_var_name = op_desc.Input("Filter").front();
+  TensorRTEngine::Weight weight;
+  if (engine->use_explicit_quantization()) {
+    auto* filter = engine->GetITensor(filter_var_name);
+    PADDLE_ENFORCE_NOT_NULL(
+        filter,
+        platform::errors::NotFound("Can not find %s ITensor in engine",
+                                   filter_var_name));
+    auto filter_dims = filter->getDimensions();
+    PADDLE_ENFORCE_EQ(
+        filter_dims.nbDims,
+        4UL,
+        platform::errors::InvalidArgument(
+            "The conv2d filter's dims size should be 4, but got %d",
+            filter_dims.nbDims));
+    n_output = filter_dims.d[0];
+    n_input = filter_dims.d[1];
+    filter_h = filter_dims.d[2];
+    filter_w = filter_dims.d[3];
+  } else {
+    auto* Y_v = scope.FindVar(filter_var_name);
+    PADDLE_ENFORCE_NOT_NULL(
+        Y_v,
+        platform::errors::NotFound("Can not find %s presistale var in scope.",
+                                   filter_var_name));
+    auto* Y_t = Y_v->GetMutable<phi::DenseTensor>();
     PADDLE_ENFORCE_EQ(
         Y_t->dims().size(),
         4UL,
@@ -63,27 +92,7 @@ void ConvertConv2d(TensorRTEngine* engine,
     n_input = Y_t->dims()[1];
     filter_h = Y_t->dims()[2];
     filter_w = Y_t->dims()[3];
-  } else {
-    filter = engine->GetITensor(op_desc.Input("Filter").front());
-    PADDLE_ENFORCE_EQ(
-        filter->getDimensions().nbDims,
-        4UL,
-        platform::errors::InvalidArgument(
-            "The conv2d filter's dims size should be 4, but got %d",
-            filter->getDimensions().nbDims));
-    n_output = filter->getDimensions().d[0];
-    n_input = filter->getDimensions().d[1];
-    filter_h = filter->getDimensions().d[2];
-    filter_w = filter->getDimensions().d[3];
-  }
-
-  bool enable_int8 = op_desc.HasAttr("enable_int8");
-
-  if (enable_int8) {
-#if IS_TRT_VERSION_GE(5000)
-    float in_scale = PADDLE_GET_CONST(float, op_desc.GetAttr("Input_scale"));
-    engine->SetTensorDynamicRange(X, in_scale);
-#endif
+    weight = engine->GetTrtWeight(op_desc.Input("Filter").front(), *Y_t);
   }
   const int groups = PADDLE_GET_CONST(int, op_desc.GetAttr("groups"));
   const std::vector<int> dilations =
@@ -124,10 +133,7 @@ void ConvertConv2d(TensorRTEngine* engine,
     nv_post_paddings.d[0] = paddings[1];
     nv_post_paddings.d[1] = paddings[3];
   }
-  TensorRTEngine::Weight weight(nvinfer1::DataType::kFLOAT, nullptr, 0);
-  if (Y_v) {
-    weight = engine->GetTrtWeight(op_desc.Input("Filter").front(), *Y_t);
-  }
+
   TensorRTEngine::Weight bias;
   bias.SetDataType(weight.get().type);
   bias.SetCount(0);
@@ -161,10 +167,7 @@ void ConvertConv2d(TensorRTEngine* engine,
   layer->setStrideNd(nv_strides);
 
   layer->setPrePadding(nv_pre_paddings);
-
-  if (!Y_v) layer->setInput(1, *filter);
-
-  if (!output_padding.empty()) {
+  if (output_padding.size() > 0) {
     nv_post_paddings.d[0] -= output_padding[0];
     nv_post_paddings.d[1] -= output_padding[1];
   }
@@ -183,6 +186,11 @@ void ConvertConv2d(TensorRTEngine* engine,
   // set dilations
   fset_dilation(layer, nv_dilations);
 
+  if (engine->use_explicit_quantization()) {
+    auto* filter_tensor = engine->GetITensor(op_desc.Input("Filter").front());
+    layer->setInput(1, *filter_tensor);
+  }
+
   auto output_name = op_desc.Output("Output").front();
   layer->setName((name + " (Output: " + output_name + ")").c_str());
   layer->getOutput(0)->setName(output_name.c_str());
diff --git a/paddle/fluid/inference/tensorrt/convert/dequantize_linear_op.cc b/paddle/fluid/inference/tensorrt/convert/dequantize_linear_op.cc
@@ -20,7 +20,7 @@ class DequantizeLinearOpConverter : public OpConverter {
   void operator()(const framework::proto::OpDesc& op,
                   const framework::Scope& scope,
                   bool test_model) override {
-#if IS_TRT_VERSION_GE(8000)
+#if IS_TRT_VERSION_GE(8510)
     VLOG(4) << "convert a dequantize_linear op to tensorrt IDequantizeLayer";
 
     // Declare inputs and attributes
@@ -45,14 +45,16 @@ class DequantizeLinearOpConverter : public OpConverter {
 
     // Add dequantize layer
     auto* layer = TRT_ENGINE_ADD_LAYER(engine_, Dequantize, *x, *scale);
-    layer->setAxis(axis);
+    if (axis >= 0) {
+      layer->setAxis(axis);
+    }
     auto output_name = op_desc.Output("Y")[0];
     RreplenishLayerAndOutput(
         layer, "dequantize_linear", {output_name}, test_model);
 #else
     PADDLE_THROW(
         platform::errors::Fatal("Paddle-TRT explicit quantization does not "
-                                "support Paddle compiled with TRT < 8.0"));
+                                "support Paddle compiled with TRT < 8.5"));
 #endif
   }
 };
diff --git a/paddle/fluid/inference/tensorrt/convert/quantize_linear_op.cc b/paddle/fluid/inference/tensorrt/convert/quantize_linear_op.cc
@@ -21,7 +21,7 @@ class QuantizeLinearOpConverter : public OpConverter {
   void operator()(const framework::proto::OpDesc& op,
                   const framework::Scope& scope,
                   bool test_model) override {
-#if IS_TRT_VERSION_GE(8000)
+#if IS_TRT_VERSION_GE(8510)
     VLOG(4) << "convert a quantize_linear op to tensorrt IQuantizeLayer";
 
     // Declare inputs and attributes
@@ -46,14 +46,16 @@ class QuantizeLinearOpConverter : public OpConverter {
 
     // Add quantize layer
     auto* layer = TRT_ENGINE_ADD_LAYER(engine_, Quantize, *x, *scale);
-    layer->setAxis(axis);
+    if (axis >= 0) {
+      layer->setAxis(axis);
+    }
     auto output_name = op_desc.Output("Y")[0];
     RreplenishLayerAndOutput(
         layer, "quantize_linear", {output_name}, test_model);
 #else
     PADDLE_THROW(
         platform::errors::Fatal("Paddle-TRT explicit quantization does not "
-                                "support Paddle compiled with TRT < 8.0"));
+                                "support Paddle compiled with TRT < 8.5"));
 #endif
   }
 };
diff --git a/paddle/fluid/inference/tensorrt/engine.cc b/paddle/fluid/inference/tensorrt/engine.cc
@@ -225,7 +225,7 @@ void TensorRTEngine::FreezeNetwork() {
 
     if (params_.calibrator) {
       infer_builder_config_->setInt8Calibrator(params_.calibrator);
-    } else {
+    } else if (!params_.use_explicit_quantization) {
       infer_builder_config_->setInt8Calibrator(nullptr);
 
       for (auto &quant_range : quant_dynamic_range_) {
diff --git a/paddle/fluid/inference/tensorrt/engine.h b/paddle/fluid/inference/tensorrt/engine.h
@@ -169,12 +169,13 @@ class TensorRTEngine {
     // affected, closing the plugin fp16 may bring some improvement on accuracy.
     bool disable_trt_plugin_fp16{false};
     int optimization_level{3};
+    bool use_explicit_quantization{false};
   };
 
   // Weight is model parameter.
   class Weight {
    public:
-    Weight() = default;
+    Weight() { w_ = nvinfer1::Weights{}; }
     Weight(nvinfer1::DataType dtype, void* value, size_t num_elem) {
       w_.type = dtype;
       w_.values = value;
@@ -527,6 +528,10 @@ class TensorRTEngine {
 
   bool LowPrecisionIOEnabled() const { return params_.enable_low_precision_io; }
 
+  bool use_explicit_quantization() const {
+    return params_.use_explicit_quantization;
+  }
+
  private:
   // Each ICudaEngine object is bound to a specific GPU when it is instantiated,
   // ensure that the thread is associated with the correct device by calling
diff --git a/paddle/fluid/inference/tensorrt/op_teller.cc b/paddle/fluid/inference/tensorrt/op_teller.cc
@@ -86,7 +86,8 @@ struct SimpleOpTypeSetTeller : public Teller {
 
   bool operator()(const framework::OpDesc& desc,
                   bool use_no_calib_int8 = false,
-                  bool with_dynamic_shape = false) override {
+                  bool with_dynamic_shape = false,
+                  bool use_explicit_quantization = false) override {
     const std::string op_type = desc.Type();
 
     std::unordered_set<std::string> control_set = {"conditional_block",
@@ -323,7 +324,7 @@ struct SimpleOpTypeSetTeller : public Teller {
       auto* block = desc.Block();
       if (block) {
         auto* filter_var_desc = block->FindVar(desc.Input("Filter")[0]);
-        if (!filter_var_desc->Persistable()) {
+        if (!filter_var_desc->Persistable() && !use_explicit_quantization) {
 #if IS_TRT_VERSION_GE(8600)
 #else
           LOG(INFO)
@@ -3093,7 +3094,8 @@ struct GenericPluginTeller : public Teller {
   GenericPluginTeller() = default;
   bool operator()(const framework::OpDesc& desc,
                   bool use_no_calib_int8 = false,
-                  bool with_dynamic_shape = false) override {
+                  bool with_dynamic_shape = false,
+                  bool use_explicit_quantization = false) override {
     const std::string op_type = desc.Type();
     // only consider dynamic_shape mode
     if (!with_dynamic_shape) {
@@ -3135,7 +3137,8 @@ struct CustomPluginTeller : public Teller {
   CustomPluginTeller() = default;
   bool operator()(const framework::OpDesc& desc,
                   bool use_no_calib_int8 = false,
-                  bool with_dynamic_shape = false) override {
+                  bool with_dynamic_shape = false,
+                  bool use_explicit_quantization = false) override {
     const std::string op_type = desc.Type();
     std::string expect_plugin_name;
 
@@ -3158,7 +3161,8 @@ struct CustomPluginTeller : public Teller {
 
 bool OpTeller::Tell(const framework::ir::Node* node,
                     bool use_no_calib_int8,
-                    bool with_dynamic_shape) {
+                    bool with_dynamic_shape,
+                    bool use_explicit_quantization) {
   const std::string op_type = node->Op()->Type();
   const framework::OpDesc desc = *node->Op();
   // do not support the op which is labeled the `skip_quant`
@@ -3168,17 +3172,26 @@ bool OpTeller::Tell(const framework::ir::Node* node,
       desc.HasAttr("skip_quant"))
     return false;
   auto& default_teller = GetDefaultTeller();
-  if ((*default_teller)(desc, use_no_calib_int8, with_dynamic_shape)) {
+  if ((*default_teller)(desc,
+                        use_no_calib_int8,
+                        with_dynamic_shape,
+                        use_explicit_quantization)) {
     SetOpConverterType(node->Op(), OpConverterType::Default);
     return true;
   }
   auto& generic_plugin_teller = GetGenericPluginTeller();
-  if ((*generic_plugin_teller)(desc, use_no_calib_int8, with_dynamic_shape)) {
+  if ((*generic_plugin_teller)(desc,
+                               use_no_calib_int8,
+                               with_dynamic_shape,
+                               use_explicit_quantization)) {
     SetOpConverterType(node->Op(), OpConverterType::GenericPluginCreater);
     return true;
   }
   auto& custom_plugin_teller = GetCustomPluginTeller();
-  if ((*custom_plugin_teller)(desc, use_no_calib_int8, with_dynamic_shape)) {
+  if ((*custom_plugin_teller)(desc,
+                              use_no_calib_int8,
+                              with_dynamic_shape,
+                              use_explicit_quantization)) {
     SetOpConverterType(node->Op(), OpConverterType::CustomPluginCreater);
     return true;
   }
diff --git a/paddle/fluid/inference/tensorrt/op_teller.h b/paddle/fluid/inference/tensorrt/op_teller.h
@@ -40,7 +40,8 @@ namespace tensorrt {
 struct Teller {
   virtual bool operator()(const framework::OpDesc& desc,
                           bool use_no_calib_int8 = false,
-                          bool with_dynamic_shape = false) = 0;
+                          bool with_dynamic_shape = false,
+                          bool use_explicit_quantization = false) = 0;
 
   virtual ~Teller() = default;
 };
@@ -74,7 +75,8 @@ class OpTeller {
 
   bool Tell(const framework::ir::Node* node,
             bool use_no_calib_int8 = false,
-            bool with_dynamic_shape = false);
+            bool with_dynamic_shape = false,
+            bool use_explicit_quantization = false);
 
   std::unique_ptr<Teller>& GetDefaultTeller() { return tellers_.at(0); }
 
diff --git a/test/ir/inference/CMakeLists.txt b/test/ir/inference/CMakeLists.txt
@@ -46,6 +46,10 @@ if(WIN32)
        "test_split_layernorm_to_math_ops_pass")
   list(REMOVE_ITEM TEST_TRT_CONVERTER
        "test_trt_convert_quantize_dequantize_linear")
+  list(REMOVE_ITEM TEST_TRT_IR_PASSES "test_trt_explicit_quantization")
+  list(REMOVE_ITEM TEST_TRT_IR_PASSES "test_trt_explicit_quantization_resnet")
+  list(REMOVE_ITEM TEST_TRT_IR_PASSES
+       "test_trt_explicit_quantization_mobilenet")
 endif()
 
 # Only for cpu(mkl + openblas)
@@ -177,6 +181,10 @@ if(WITH_GPU AND TENSORRT_FOUND)
   set_tests_properties(test_trt_inference_predictor PROPERTIES TIMEOUT 60)
   set_tests_properties(test_trt_inference_fp16_io PROPERTIES TIMEOUT 300)
   set_tests_properties(test_trt_optimization_level PROPERTIES TIMEOUT 300)
+  set_tests_properties(test_trt_explicit_quantization_resnet PROPERTIES TIMEOUT
+                                                                        300)
+  set_tests_properties(test_trt_explicit_quantization_mobilenet
+                       PROPERTIES TIMEOUT 300)
   if(WITH_MKLDNN)
     set_tests_properties(test_save_optimized_model_pass PROPERTIES TIMEOUT 300)
   endif()
diff --git a/test/ir/inference/test_trt_explicit_quantization.py b/test/ir/inference/test_trt_explicit_quantization.py
diff --git a/test/ir/inference/test_trt_explicit_quantization_mobilenet.py b/test/ir/inference/test_trt_explicit_quantization_mobilenet.py
diff --git a/test/ir/inference/test_trt_explicit_quantization_model.py b/test/ir/inference/test_trt_explicit_quantization_model.py
diff --git a/test/ir/inference/test_trt_explicit_quantization_resnet.py b/test/ir/inference/test_trt_explicit_quantization_resnet.py
