未验证 提交 994673bf 编写于 作者: P Pei Yang 提交者: GitHub

change avg pooling and global pooling to trt layer in dynamic shape mode (#28702)

* change avg pooling and global pooling to trt layer

* add support for static shape global pooling

* modify trt errmsg
上级 3416c359
...@@ -97,13 +97,17 @@ class Pool2dOpConverter : public OpConverter { ...@@ -97,13 +97,17 @@ class Pool2dOpConverter : public OpConverter {
adaptive = BOOST_GET_CONST(bool, op_desc.GetAttr("adaptive")); adaptive = BOOST_GET_CONST(bool, op_desc.GetAttr("adaptive"));
nvinfer1::PoolingType nv_pool_type = nvinfer1::PoolingType::kMAX; nvinfer1::PoolingType nv_pool_type = nvinfer1::PoolingType::kMAX;
nvinfer1::ReduceOperation reduce_operation =
nvinfer1::ReduceOperation::kMAX;
plugin::PoolPlugin::PoolType plugin_pool_type = plugin::PoolPlugin::PoolType plugin_pool_type =
plugin::PoolPlugin::PoolType::max; plugin::PoolPlugin::PoolType::max;
if (pool_type == "max") { if (pool_type == "max") {
nv_pool_type = nvinfer1::PoolingType::kMAX; nv_pool_type = nvinfer1::PoolingType::kMAX;
reduce_operation = nvinfer1::ReduceOperation::kMAX;
plugin_pool_type = plugin::PoolPlugin::PoolType::max; plugin_pool_type = plugin::PoolPlugin::PoolType::max;
} else if (pool_type == "avg") { } else if (pool_type == "avg") {
nv_pool_type = nvinfer1::PoolingType::kAVERAGE; nv_pool_type = nvinfer1::PoolingType::kAVERAGE;
reduce_operation = nvinfer1::ReduceOperation::kAVG;
plugin_pool_type = plugin::PoolPlugin::PoolType::avg; plugin_pool_type = plugin::PoolPlugin::PoolType::avg;
} else { } else {
PADDLE_THROW(platform::errors::Fatal( PADDLE_THROW(platform::errors::Fatal(
...@@ -126,12 +130,17 @@ class Pool2dOpConverter : public OpConverter { ...@@ -126,12 +130,17 @@ class Pool2dOpConverter : public OpConverter {
} }
if (engine_->with_dynamic_shape()) { 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, auto *pool_layer = TRT_ENGINE_ADD_LAYER(engine_, Pooling, *input1,
nv_pool_type, nv_ksize); nv_pool_type, nv_ksize);
pool_layer->setStride(nv_strides); pool_layer->setStride(nv_strides);
pool_layer->setPadding(nv_paddings); pool_layer->setPadding(nv_paddings);
pool_layer->setAverageCountExcludesPadding(exclusive);
layer = pool_layer; 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 { } else {
#if IS_TRT_VERSION_GE(6000) #if IS_TRT_VERSION_GE(6000)
plugin::PoolPluginDynamic *plugin = plugin::PoolPluginDynamic *plugin =
...@@ -153,16 +162,20 @@ class Pool2dOpConverter : public OpConverter { ...@@ -153,16 +162,20 @@ class Pool2dOpConverter : public OpConverter {
if (global_pooling == true) { if (global_pooling == true) {
nv_ksize.d[0] = input_shape.d[input_dims - 2]; nv_ksize.d[0] = input_shape.d[input_dims - 2];
nv_ksize.d[1] = input_shape.d[input_dims - 1]; 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<nvinfer1::ITensor *>(input1), engine_, Pooling, *const_cast<nvinfer1::ITensor *>(input1),
nv_pool_type, nv_ksize); nv_pool_type, nv_ksize);
PADDLE_ENFORCE_NOT_NULL( PADDLE_ENFORCE_NOT_NULL(
layer, platform::errors::Fatal( pool_layer, platform::errors::Fatal(
"trt pool layer in converter could not be created.")); "trt pool layer in converter could not be created."));
auto output_name = op_desc.Output("Out")[0]; auto output_name = op_desc.Output("Out")[0];
layer->setName(("pool2d (Output: " + output_name + ")").c_str()); pool_layer->setStride(nv_strides);
layer->getOutput(0)->setName(output_name.c_str()); pool_layer->setPadding(nv_paddings);
engine_->SetITensor(output_name, layer->getOutput(0)); 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) { if (test_mode) {
engine_->DeclareOutput(output_name); engine_->DeclareOutput(output_name);
} }
......
...@@ -42,7 +42,7 @@ void* GetDsoHandle(const std::string& dso_name) { ...@@ -42,7 +42,7 @@ void* GetDsoHandle(const std::string& dso_name) {
if (nullptr == dso_handle) { if (nullptr == dso_handle) {
auto error_msg = auto error_msg =
"You are using Paddle compiled with TensorRT, but TensorRT dynamic " "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; std::cerr << error_msg;
} }
return dso_handle; return dso_handle;
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册