From 32211fe9c4c22168dfb73f19763b17ac9191341a Mon Sep 17 00:00:00 2001
From: Pei Yang <peiyang@baidu.com>
Date: Wed, 3 Mar 2021 15:17:55 +0800
Subject: [PATCH] TRT conv2d converter support SAME padding (#31379)

---
 paddle/fluid/inference/tensorrt/convert/conv2d_op.cc     | 7 +++++++
 paddle/fluid/inference/tensorrt/op_teller.cc             | 8 +-------
 .../tests/unittests/ir/inference/test_trt_conv_pass.py   | 9 ++-------
 3 files changed, 10 insertions(+), 14 deletions(-)

diff --git a/paddle/fluid/inference/tensorrt/convert/conv2d_op.cc b/paddle/fluid/inference/tensorrt/convert/conv2d_op.cc
index 6871d53f42..5515cd35da 100644
--- a/paddle/fluid/inference/tensorrt/convert/conv2d_op.cc
+++ b/paddle/fluid/inference/tensorrt/convert/conv2d_op.cc
@@ -97,6 +97,10 @@ void ConvertConv2d(TensorRTEngine* engine, const framework::proto::OpDesc& op,
       BOOST_GET_CONST(std::vector<int>, op_desc.GetAttr("strides"));
   const std::vector<int> paddings =
       BOOST_GET_CONST(std::vector<int>, op_desc.GetAttr("paddings"));
+  std::string padding_algorithm = "EXPLICIT";
+  if (op_desc.HasAttr("padding_algorithm"))
+    padding_algorithm =
+        BOOST_GET_CONST(std::string, op_desc.GetAttr("padding_algorithm"));
 
   nvinfer1::DimsHW nv_ksize(filter_h, filter_w);
   nvinfer1::DimsHW nv_dilations(dilations[0], dilations[1]);
@@ -126,6 +130,9 @@ void ConvertConv2d(TensorRTEngine* engine, const framework::proto::OpDesc& op,
   layer->setStride(nv_strides);
   layer->setPadding(nv_paddings);
   layer->setNbGroups(groups);
+  if (padding_algorithm == "SAME") {
+    layer->setPaddingMode(nvinfer1::PaddingMode::kSAME_UPPER);
+  }
   // set dilations
   fset_dilation(layer, nv_dilations);
 
diff --git a/paddle/fluid/inference/tensorrt/op_teller.cc b/paddle/fluid/inference/tensorrt/op_teller.cc
index 6f000fbccf..052d17878a 100644
--- a/paddle/fluid/inference/tensorrt/op_teller.cc
+++ b/paddle/fluid/inference/tensorrt/op_teller.cc
@@ -129,13 +129,7 @@ bool OpTeller::Tell(const framework::ir::Node* node, bool use_no_calib_int8,
       std::vector<int> paddings =
           BOOST_GET_CONST(std::vector<int>, desc.GetAttr("paddings"));
 
-      std::string padding_algorithm = "EXPLICIT";
-      if (desc.HasAttr("padding_algorithm"))
-        padding_algorithm =
-            BOOST_GET_CONST(std::string, desc.GetAttr("padding_algorithm"));
-      if (paddings.size() > 2 ||
-          (padding_algorithm == "SAME" && op_type != "pool2d"))
-        return false;
+      if (paddings.size() > 2) return false;
     }
     if (op_type == "matmul") {
       auto* block = desc.Block();
diff --git a/python/paddle/fluid/tests/unittests/ir/inference/test_trt_conv_pass.py b/python/paddle/fluid/tests/unittests/ir/inference/test_trt_conv_pass.py
index 0de37fce0a..0821b390e5 100644
--- a/python/paddle/fluid/tests/unittests/ir/inference/test_trt_conv_pass.py
+++ b/python/paddle/fluid/tests/unittests/ir/inference/test_trt_conv_pass.py
@@ -67,15 +67,12 @@ class TensorRTSubgraphPassConvValidPaddingTest(TensorRTSubgraphPassConvTest):
         self.conv_padding = 'VALID'
 
 
-'''
-# conv2d padded in 'SAME' mode is not yet supported in TRT, reopen this when support is complete.
 class TensorRTSubgraphPassConvSamePaddingTest(InferencePassTest):
     def set_params(self):
         self.conv_num_filters = 6
         self.conv_filter_size = 6
         self.conv_groups = 3
         self.conv_padding = 'SAME'
-'''
 
 
 class TensorRTSubgraphPassDepthwiseConvTest(TensorRTSubgraphPassConvTest):
@@ -131,15 +128,13 @@ class TensorRTSubgraphPassConvTransposeValidPaddingTest(
         self.conv_padding = 'VALID'
 
 
-'''
-# conv2d_transpose padded in 'SAME' mode is not yet supported in TRT, reopen this when support is complete.
-class TensorRTSubgraphPassConvTransposeSamePaddingTest(TensorRTSubgraphPassConvTransposeTest):
+class TensorRTSubgraphPassConvTransposeSamePaddingTest(
+        TensorRTSubgraphPassConvTransposeTest):
     def set_params(self):
         self.conv_num_filters = 6
         self.conv_filter_size = 6
         self.conv_groups = 1
         self.conv_padding = 'SAME'
-'''
 
 
 class TensorRTSubgraphPassDepthwiseConvTransposeTest(
-- 
GitLab