From 602d2ce5c9eac603dc74e04a5695aa1894f8decf Mon Sep 17 00:00:00 2001 From: Pei Yang Date: Wed, 21 Oct 2020 13:13:08 +0800 Subject: [PATCH] change avg pooling from trt plugin to trt layer (#28032) --- paddle/fluid/inference/tensorrt/convert/pool2d_op.cc | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/paddle/fluid/inference/tensorrt/convert/pool2d_op.cc b/paddle/fluid/inference/tensorrt/convert/pool2d_op.cc index c031630f36d..303130e74f5 100644 --- a/paddle/fluid/inference/tensorrt/convert/pool2d_op.cc +++ b/paddle/fluid/inference/tensorrt/convert/pool2d_op.cc @@ -88,6 +88,9 @@ class Pool2dOpConverter : public OpConverter { BOOST_GET_CONST(std::vector, op_desc.GetAttr("strides")); std::vector paddings = BOOST_GET_CONST(std::vector, op_desc.GetAttr("paddings")); + bool exclusive = op_desc.HasAttr("exclusive") + ? BOOST_GET_CONST(bool, op_desc.GetAttr("exclusive")) + : true; bool ceil_mode = BOOST_GET_CONST(bool, op_desc.GetAttr("ceil_mode")); bool adaptive = false; if (op_desc.HasAttr("adaptive")) @@ -166,7 +169,7 @@ class Pool2dOpConverter : public OpConverter { return; } - if (!adaptive && pool_type == "max") { + 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. @@ -194,6 +197,7 @@ class Pool2dOpConverter : public OpConverter { "trt pool layer in converter could not be created.")); pool_layer->setStride(nv_strides); pool_layer->setPadding(nv_paddings); + pool_layer->setAverageCountExcludesPadding(exclusive); layer = pool_layer; } else { // Average pooling needs to exclude the padding pixels from the average @@ -213,7 +217,6 @@ class Pool2dOpConverter : public OpConverter { "trt pool plugin layer in converter could not be created.")); layer = pool_layer; } - auto output_name = op_desc.Output("Out")[0]; RreplenishLayerAndOutput(layer, "pool2d", {output_name}, test_mode); } -- GitLab