From 192e08cbff30ff2d602aec85ef1bf5b3252590e6 Mon Sep 17 00:00:00 2001
From: wenbin <wang3323032@qq.com>
Date: Wed, 13 Oct 2021 14:53:07 +0800
Subject: [PATCH] pool fix (#36388)

* pool fix

* comments
---
 .../inference/tensorrt/convert/pool2d_op.cc   | 48 +++++++++++++------
 1 file changed, 33 insertions(+), 15 deletions(-)

diff --git a/paddle/fluid/inference/tensorrt/convert/pool2d_op.cc b/paddle/fluid/inference/tensorrt/convert/pool2d_op.cc
index 733a8f64ae5..e03842db2b8 100644
--- a/paddle/fluid/inference/tensorrt/convert/pool2d_op.cc
+++ b/paddle/fluid/inference/tensorrt/convert/pool2d_op.cc
@@ -115,6 +115,18 @@ class Pool2dOpConverter : public OpConverter {
     nvinfer1::DimsHW nv_paddings(paddings[0], paddings[1]);
 
     nvinfer1::ILayer *layer = nullptr;
+    nvinfer1::DimsHW pre_pad(0, 0);
+    nvinfer1::DimsHW post_pad(0, 0);
+    // paddle Non ceil_mode : Output size = (input size - filter size + 2 *
+    // padding) / stride (stride size) + 1
+    // tensorrt EXPLICIT_ROUND_DOWN: O = floor((M - DK) / S) + 1
+    // so if M - DK < 0 we need extra padding
+    if (input_shape.d[input_dims - 2] - ksize[0] + 2 * paddings[0] < 0) {
+      post_pad.h() = strides[0] - 1;
+    }
+    if (input_shape.d[input_dims - 1] - ksize[1] + 2 * paddings[1] < 0) {
+      post_pad.w() = strides[1] - 1;
+    }
 
     if (op_desc.HasAttr("enable_int8")) {
 #if IS_TRT_VERSION_GE(5000)
@@ -126,6 +138,16 @@ class Pool2dOpConverter : public OpConverter {
 
     if (engine_->with_dynamic_shape()) {
       if (!adaptive && !global_pooling && !ceil_mode) {
+        if ((post_pad.w() > 0 || post_pad.h() > 0) &&
+            (padding_algorithm != "SAME")) {
+          auto *pad_layer = TRT_ENGINE_ADD_LAYER(engine_, Padding, *input1,
+                                                 pre_pad, post_pad);
+          PADDLE_ENFORCE_NOT_NULL(
+              pad_layer, platform::errors::Fatal(
+                             "Pad layer in poolOp converter could not be "
+                             "created. The pointer to pad layer is `NULL`."));
+          input1 = pad_layer->getOutput(0);
+        }
         auto *pool_layer = TRT_ENGINE_ADD_LAYER(engine_, Pooling, *input1,
                                                 nv_pool_type, nv_ksize);
         pool_layer->setStride(nv_strides);
@@ -184,9 +206,8 @@ class Pool2dOpConverter : public OpConverter {
     if (global_pooling == true) {
       nv_ksize.d[0] = input_shape.d[input_dims - 2];
       nv_ksize.d[1] = input_shape.d[input_dims - 1];
-      auto *pool_layer = TRT_ENGINE_ADD_LAYER(
-          engine_, Pooling, *const_cast<nvinfer1::ITensor *>(input1),
-          nv_pool_type, nv_ksize);
+      auto *pool_layer = TRT_ENGINE_ADD_LAYER(engine_, Pooling, *input1,
+                                              nv_pool_type, nv_ksize);
       PADDLE_ENFORCE_NOT_NULL(
           pool_layer, platform::errors::Fatal(
                           "trt pool layer in converter could not be created."));
@@ -208,28 +229,25 @@ class Pool2dOpConverter : public OpConverter {
     }
 
     if (!adaptive) {
-      // Under ceil mode, the pre_pad and post_pad are used to
-      // record the the padding size. In some ceil mode cases,
-      // we do not need padding, so we initialize the two vars to 0.
-
-      nvinfer1::DimsHW pre_pad(0, 0);
-      nvinfer1::DimsHW post_pad(0, 0);
       if (ceil_mode) {
         // If ceil mode is true, we will pad the appropriate size to the input.
         DealCeilMode(input_shape, ksize, strides, paddings, &pre_pad, &post_pad,
                      input_dims);
-        auto *pad_layer = TRT_ENGINE_ADD_LAYER(
-            engine_, Padding, *const_cast<nvinfer1::ITensor *>(input1), pre_pad,
-            post_pad);
+      }
+
+      if ((post_pad.w() > 0 || post_pad.h() > 0) &&
+          (padding_algorithm != "SAME")) {
+        auto *pad_layer =
+            TRT_ENGINE_ADD_LAYER(engine_, Padding, *input1, pre_pad, post_pad);
         PADDLE_ENFORCE_NOT_NULL(
             pad_layer, platform::errors::Fatal(
                            "Pad layer in poolOp converter could not be "
                            "created. The pointer to pad layer is `NULL`."));
         input1 = pad_layer->getOutput(0);
       }
-      auto *pool_layer = TRT_ENGINE_ADD_LAYER(
-          engine_, Pooling, *const_cast<nvinfer1::ITensor *>(input1),
-          nv_pool_type, nv_ksize);
+
+      auto *pool_layer = TRT_ENGINE_ADD_LAYER(engine_, Pooling, *input1,
+                                              nv_pool_type, nv_ksize);
       PADDLE_ENFORCE_NOT_NULL(
           pool_layer, platform::errors::Fatal(
                           "trt pool layer in converter could not be created."));
-- 
GitLab