diff --git a/paddle/fluid/inference/tensorrt/convert/pool2d_op.cc b/paddle/fluid/inference/tensorrt/convert/pool2d_op.cc index 303130e74f512abe041fd14a8bcbaea9400fa902..ca5a1a77bd0e8ee8f35fecc838ad303601661d91 100644 --- a/paddle/fluid/inference/tensorrt/convert/pool2d_op.cc +++ b/paddle/fluid/inference/tensorrt/convert/pool2d_op.cc @@ -97,13 +97,17 @@ class Pool2dOpConverter : public OpConverter { adaptive = BOOST_GET_CONST(bool, op_desc.GetAttr("adaptive")); nvinfer1::PoolingType nv_pool_type = nvinfer1::PoolingType::kMAX; + nvinfer1::ReduceOperation reduce_operation = + nvinfer1::ReduceOperation::kMAX; plugin::PoolPlugin::PoolType plugin_pool_type = plugin::PoolPlugin::PoolType::max; if (pool_type == "max") { nv_pool_type = nvinfer1::PoolingType::kMAX; + reduce_operation = nvinfer1::ReduceOperation::kMAX; plugin_pool_type = plugin::PoolPlugin::PoolType::max; } else if (pool_type == "avg") { nv_pool_type = nvinfer1::PoolingType::kAVERAGE; + reduce_operation = nvinfer1::ReduceOperation::kAVG; plugin_pool_type = plugin::PoolPlugin::PoolType::avg; } else { PADDLE_THROW(platform::errors::Fatal( @@ -126,12 +130,17 @@ class Pool2dOpConverter : public OpConverter { } if (engine_->with_dynamic_shape()) { - if (!adaptive && pool_type == "max" && !global_pooling && !ceil_mode) { + if (!adaptive && !global_pooling && !ceil_mode) { auto *pool_layer = TRT_ENGINE_ADD_LAYER(engine_, Pooling, *input1, nv_pool_type, nv_ksize); pool_layer->setStride(nv_strides); pool_layer->setPadding(nv_paddings); + pool_layer->setAverageCountExcludesPadding(exclusive); layer = pool_layer; + } else if (global_pooling) { + auto *reduce_layer = TRT_ENGINE_ADD_LAYER(engine_, Reduce, *input1, + reduce_operation, 12, true); + layer = reduce_layer; } else { #if IS_TRT_VERSION_GE(6000) plugin::PoolPluginDynamic *plugin = @@ -153,16 +162,20 @@ 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 *layer = TRT_ENGINE_ADD_LAYER( + auto *pool_layer = TRT_ENGINE_ADD_LAYER( engine_, Pooling, *const_cast(input1), nv_pool_type, nv_ksize); PADDLE_ENFORCE_NOT_NULL( - layer, platform::errors::Fatal( - "trt pool layer in converter could not be created.")); + pool_layer, platform::errors::Fatal( + "trt pool layer in converter could not be created.")); auto output_name = op_desc.Output("Out")[0]; - layer->setName(("pool2d (Output: " + output_name + ")").c_str()); - layer->getOutput(0)->setName(output_name.c_str()); - engine_->SetITensor(output_name, layer->getOutput(0)); + pool_layer->setStride(nv_strides); + pool_layer->setPadding(nv_paddings); + pool_layer->setAverageCountExcludesPadding(exclusive); + pool_layer->setName(("pool2d (Output: " + output_name + ")").c_str()); + pool_layer->getOutput(0)->setName(output_name.c_str()); + engine_->SetITensor(output_name, pool_layer->getOutput(0)); + layer = pool_layer; if (test_mode) { engine_->DeclareOutput(output_name); } diff --git a/paddle/fluid/platform/dynload/tensorrt.cc b/paddle/fluid/platform/dynload/tensorrt.cc index 6232a6e33cac453af92ffec26c94729e58b3904e..e72fbd246cf05ddc4f8e9455ba9c579835faa5bd 100644 --- a/paddle/fluid/platform/dynload/tensorrt.cc +++ b/paddle/fluid/platform/dynload/tensorrt.cc @@ -42,7 +42,7 @@ void* GetDsoHandle(const std::string& dso_name) { if (nullptr == dso_handle) { auto error_msg = "You are using Paddle compiled with TensorRT, but TensorRT dynamic " - "library is not found. Ignore this if TensorRT is not needed."; + "library is not found. Ignore this if TensorRT is not needed.\n"; std::cerr << error_msg; } return dso_handle;