未验证 提交 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 {
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<nvinfer1::ITensor *>(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);
}
......
......@@ -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;
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册