From d91352c040c965b04e7c3e42f9ccb71291bb545e Mon Sep 17 00:00:00 2001 From: Pei Yang Date: Mon, 28 Jun 2021 10:39:57 +0800 Subject: [PATCH] [Paddle-TRT]Fix flatten converter when batch_size > 1 (#33768) * fix trt flatten converter when batch_size > 1 * change ut to same dynamic shape --- .../fluid/inference/tensorrt/convert/flatten_op.cc | 13 +++++++++++-- paddle/fluid/inference/tensorrt/helper.h | 13 +++++++++++++ .../unittests/ir/inference/test_trt_flatten_op.py | 8 ++++---- 3 files changed, 28 insertions(+), 6 deletions(-) diff --git a/paddle/fluid/inference/tensorrt/convert/flatten_op.cc b/paddle/fluid/inference/tensorrt/convert/flatten_op.cc index 322b42667fa..25351cc10ec 100644 --- a/paddle/fluid/inference/tensorrt/convert/flatten_op.cc +++ b/paddle/fluid/inference/tensorrt/convert/flatten_op.cc @@ -53,10 +53,19 @@ class FlattenOpConverter : public OpConverter { layer->setReshapeDimensions(flatten_dim); } else { auto* shape_layer = TRT_ENGINE_ADD_LAYER(engine_, Shape, *input); + nvinfer1::Dims start_dim, size_dim, stride_dim; + start_dim.nbDims = 1; + size_dim.nbDims = 1; + stride_dim.nbDims = 1; + start_dim.d[0] = 1; + size_dim.d[0] = dims - 1; + stride_dim.d[0] = 1; + auto* slice_layer = + TRT_ENGINE_ADD_LAYER(engine_, Slice, *(shape_layer->getOutput(0)), + start_dim, size_dim, stride_dim); uint32_t reduce_dim = 1; - auto* reduce_prod_layer = TRT_ENGINE_ADD_LAYER( - engine_, Reduce, *(shape_layer->getOutput(0)), + engine_, Reduce, *(slice_layer->getOutput(0)), nvinfer1::ReduceOperation::kPROD, reduce_dim, true); int32_t* constant_weight_data = new int32_t[1]; constant_weight_data[0] = -1; diff --git a/paddle/fluid/inference/tensorrt/helper.h b/paddle/fluid/inference/tensorrt/helper.h index 8b557c6e6b8..e3c7d8b1033 100644 --- a/paddle/fluid/inference/tensorrt/helper.h +++ b/paddle/fluid/inference/tensorrt/helper.h @@ -134,6 +134,19 @@ inline size_t ProductDim(const nvinfer1::Dims& dims) { return v; } +inline void PrintITensorShape(nvinfer1::ITensor* X) { + auto dims = X->getDimensions(); + auto name = X->getName(); + std::cout << "ITensor " << name << " shape: ["; + for (int i = 0; i < dims.nbDims; i++) { + if (i == dims.nbDims - 1) + std::cout << dims.d[i]; + else + std::cout << dims.d[i] << ", "; + } + std::cout << "]\n"; +} + } // namespace tensorrt } // namespace inference } // namespace paddle diff --git a/python/paddle/fluid/tests/unittests/ir/inference/test_trt_flatten_op.py b/python/paddle/fluid/tests/unittests/ir/inference/test_trt_flatten_op.py index 1f8f829d27c..bb28fcf7085 100644 --- a/python/paddle/fluid/tests/unittests/ir/inference/test_trt_flatten_op.py +++ b/python/paddle/fluid/tests/unittests/ir/inference/test_trt_flatten_op.py @@ -63,10 +63,10 @@ class TRTFlattenDynamicTest(InferencePassTest): self.trt_parameters = TRTFlattenDynamicTest.TensorRTParam( 1 << 30, 32, 0, AnalysisConfig.Precision.Float32, False, False) self.dynamic_shape_params = TRTFlattenDynamicTest.DynamicShapeParam({ - 'data': [1, 6, 8, 8], - 'flatten_0.tmp_0': [1, 6 * 8 * 8] - }, {'data': [3, 6, 128, 128], - 'flatten_0.tmp_0': [3, 6 * 128 * 128]}, { + 'data': [2, 6, 64, 64], + 'flatten_0.tmp_0': [2, 6 * 64 * 64] + }, {'data': [2, 6, 64, 64], + 'flatten_0.tmp_0': [2, 6 * 64 * 64]}, { 'data': [2, 6, 64, 64], 'flatten_0.tmp_0': [2, 6 * 64 * 64] }, False) -- GitLab