[PaddleTRT] Fixed bug in generating quantization calibration table fo…

…r yolo_box (PaddlePaddle#61596) * 修复yolov3量化表yolo_box的bug
lizexu123 · Feb 5, 2024 · 0f23c62 · 0f23c62
1 parent 60325a1
commit 0f23c62
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 5 deletions.
diff --git a/paddle/fluid/inference/tensorrt/convert/yolo_box_op.cc b/paddle/fluid/inference/tensorrt/convert/yolo_box_op.cc
@@ -49,10 +49,13 @@ class YoloBoxOpConverter : public OpConverter {
             ? PADDLE_GET_CONST(float, op_desc.GetAttr("iou_aware_factor"))
             : 0.5;
 
-    int type_id = static_cast<int>(engine_->WithFp16());
+    bool with_fp16 = engine_->WithFp16() && !engine_->disable_trt_plugin_fp16();
+    if (engine_->precision() == phi::DataType::INT8) {
+      with_fp16 = true;
+    }
     auto input_dim = X_tensor->getDimensions();
     auto* yolo_box_plugin = new plugin::YoloBoxPlugin(
-        type_id ? nvinfer1::DataType::kHALF : nvinfer1::DataType::kFLOAT,
+        with_fp16 ? nvinfer1::DataType::kHALF : nvinfer1::DataType::kFLOAT,
         anchors,
         class_num,
         conf_thresh,

diff --git a/paddle/fluid/inference/tensorrt/plugin/yolo_box_op_plugin.cu b/paddle/fluid/inference/tensorrt/plugin/yolo_box_op_plugin.cu
@@ -12,11 +12,10 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include "paddle/fluid/inference/tensorrt/plugin/yolo_box_op_plugin.h"
 #include <algorithm>
 #include <cassert>
 
-#include "paddle/fluid/inference/tensorrt/plugin/yolo_box_op_plugin.h"
-
 namespace paddle {
 namespace inference {
 namespace tensorrt {
@@ -102,7 +101,9 @@ nvinfer1::Dims YoloBoxPlugin::getOutputDimensions(
 
 bool YoloBoxPlugin::supportsFormat(
     nvinfer1::DataType type, nvinfer1::TensorFormat format) const TRT_NOEXCEPT {
-  return ((type == data_type_ || type == nvinfer1::DataType::kINT32) &&
+  return ((type == nvinfer1::DataType::kFLOAT ||
+           type == nvinfer1::DataType::kHALF ||
+           type == nvinfer1::DataType::kINT32) &&
           format == nvinfer1::TensorFormat::kLINEAR);
 }