From 7ee9ba2f29f78ab669ee35f9d32e2c967e67c849 Mon Sep 17 00:00:00 2001
From: JingZhuangzhuang <75348594+JZZ-NOTE@users.noreply.github.com>
Date: Mon, 18 Apr 2022 10:59:24 +0800
Subject: [PATCH] fix_poo2d_trt_convert (#41860)

---
 .../inference/tensorrt/convert/pool2d_op.cc   | 58 ++++++++++++++-----
 1 file changed, 44 insertions(+), 14 deletions(-)

diff --git a/paddle/fluid/inference/tensorrt/convert/pool2d_op.cc b/paddle/fluid/inference/tensorrt/convert/pool2d_op.cc
index 7824a0f1e2..29acf549cb 100644
--- a/paddle/fluid/inference/tensorrt/convert/pool2d_op.cc
+++ b/paddle/fluid/inference/tensorrt/convert/pool2d_op.cc
@@ -256,19 +256,51 @@ class Pool2dOpConverter : public OpConverter {
 
     if (!adaptive) {
       if (ceil_mode) {
-        std::vector<int> input_shape_v;
-        for (int i = 0; i < input_dims; i++) {
-          input_shape_v.push_back(input_shape.d[i]);
+        if (nv_ksize.d[0] % nv_strides.d[0] == 0 &&
+            nv_ksize.d[1] % nv_strides.d[1] == 0) {
+          nvinfer1::DimsHW pre_pad(0, 0);
+          nvinfer1::DimsHW post_pad(0, 0);
+          // 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, *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);
+          PADDLE_ENFORCE_NOT_NULL(
+              pool_layer,
+              platform::errors::Fatal(
+                  "trt pool layer in converter could not be created."));
+          pool_layer->setStride(nv_strides);
+          pool_layer->setPadding(nv_paddings);
+          if (padding_algorithm == "SAME") {
+            pool_layer->setPaddingMode(nvinfer1::PaddingMode::kSAME_UPPER);
+          }
+          pool_layer->setAverageCountExcludesPadding(exclusive);
+          layer = pool_layer;
+        } else {
+          std::vector<int> input_shape_v;
+          for (int i = 0; i < input_dims; i++) {
+            input_shape_v.push_back(input_shape.d[i]);
+          }
+          plugin::PoolPlugin *plugin = new plugin::PoolPlugin(
+              ceil_mode, plugin_pool_type, adaptive, exclusive, ksize, strides,
+              paddings, input_shape_v, real_paddings);
+          auto *pool_layer = engine_->AddPlugin(&input1, 1, plugin);
+          PADDLE_ENFORCE_NOT_NULL(
+              pool_layer,
+              platform::errors::Fatal(
+                  "trt pool plugin layer in converter could not be created."));
+          layer = pool_layer;
         }
-        plugin::PoolPlugin *plugin = new plugin::PoolPlugin(
-            ceil_mode, plugin_pool_type, adaptive, exclusive, ksize, strides,
-            paddings, input_shape_v, real_paddings);
-        auto *pool_layer = engine_->AddPlugin(&input1, 1, plugin);
-        PADDLE_ENFORCE_NOT_NULL(
-            pool_layer,
-            platform::errors::Fatal(
-                "trt pool plugin layer in converter could not be created."));
-        layer = pool_layer;
       } else {
 #if IS_TRT_VERSION_GE(8000)
         // Exclude padding pixels from the average mean is not supported well by
@@ -299,7 +331,6 @@ class Pool2dOpConverter : public OpConverter {
         pool_layer->setAverageCountExcludesPadding(exclusive);
         layer = pool_layer;
       }
-
     } else {
       // Average pooling needs to exclude the padding pixels from the average
       // mean.
@@ -327,5 +358,4 @@ class Pool2dOpConverter : public OpConverter {
 }  // namespace inference
 }  // namespace paddle
 
-USE_OP_ITSELF(pool2d);
 REGISTER_TRT_OP_CONVERTER(pool2d, Pool2dOpConverter);
-- 
GitLab