未验证 提交 192e08cb 编写于 作者: W wenbin 提交者: GitHub

pool fix (#36388)

* pool fix

* comments
上级 bf748f24
......@@ -115,6 +115,18 @@ class Pool2dOpConverter : public OpConverter {
nvinfer1::DimsHW nv_paddings(paddings[0], paddings[1]);
nvinfer1::ILayer *layer = nullptr;
nvinfer1::DimsHW pre_pad(0, 0);
nvinfer1::DimsHW post_pad(0, 0);
// paddle Non ceil_mode : Output size = (input size - filter size + 2 *
// padding) / stride (stride size) + 1
// tensorrt EXPLICIT_ROUND_DOWN: O = floor((M - DK) / S) + 1
// so if M - DK < 0 we need extra padding
if (input_shape.d[input_dims - 2] - ksize[0] + 2 * paddings[0] < 0) {
post_pad.h() = strides[0] - 1;
}
if (input_shape.d[input_dims - 1] - ksize[1] + 2 * paddings[1] < 0) {
post_pad.w() = strides[1] - 1;
}
if (op_desc.HasAttr("enable_int8")) {
#if IS_TRT_VERSION_GE(5000)
......@@ -126,6 +138,16 @@ class Pool2dOpConverter : public OpConverter {
if (engine_->with_dynamic_shape()) {
if (!adaptive && !global_pooling && !ceil_mode) {
if ((post_pad.w() > 0 || post_pad.h() > 0) &&
(padding_algorithm != "SAME")) {
auto *pad_layer = TRT_ENGINE_ADD_LAYER(engine_, Padding, *input1,
pre_pad, post_pad);
PADDLE_ENFORCE_NOT_NULL(
pad_layer, platform::errors::Fatal(
"Pad layer in poolOp converter could not be "
"created. The pointer to pad layer is `NULL`."));
input1 = pad_layer->getOutput(0);
}
auto *pool_layer = TRT_ENGINE_ADD_LAYER(engine_, Pooling, *input1,
nv_pool_type, nv_ksize);
pool_layer->setStride(nv_strides);
......@@ -184,8 +206,7 @@ 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 *pool_layer = TRT_ENGINE_ADD_LAYER(
engine_, Pooling, *const_cast<nvinfer1::ITensor *>(input1),
auto *pool_layer = TRT_ENGINE_ADD_LAYER(engine_, Pooling, *input1,
nv_pool_type, nv_ksize);
PADDLE_ENFORCE_NOT_NULL(
pool_layer, platform::errors::Fatal(
......@@ -208,27 +229,24 @@ class Pool2dOpConverter : public OpConverter {
}
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.
nvinfer1::DimsHW pre_pad(0, 0);
nvinfer1::DimsHW post_pad(0, 0);
if (ceil_mode) {
// If ceil mode is true, we will pad the appropriate size to the input.
DealCeilMode(input_shape, ksize, strides, paddings, &pre_pad, &post_pad,
input_dims);
auto *pad_layer = TRT_ENGINE_ADD_LAYER(
engine_, Padding, *const_cast<nvinfer1::ITensor *>(input1), pre_pad,
post_pad);
}
if ((post_pad.w() > 0 || post_pad.h() > 0) &&
(padding_algorithm != "SAME")) {
auto *pad_layer =
TRT_ENGINE_ADD_LAYER(engine_, Padding, *input1, pre_pad, post_pad);
PADDLE_ENFORCE_NOT_NULL(
pad_layer, platform::errors::Fatal(
"Pad layer in poolOp converter could not be "
"created. The pointer to pad layer is `NULL`."));
input1 = pad_layer->getOutput(0);
}
auto *pool_layer = TRT_ENGINE_ADD_LAYER(
engine_, Pooling, *const_cast<nvinfer1::ITensor *>(input1),
auto *pool_layer = TRT_ENGINE_ADD_LAYER(engine_, Pooling, *input1,
nv_pool_type, nv_ksize);
PADDLE_ENFORCE_NOT_NULL(
pool_layer, platform::errors::Fatal(
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册