From de992400b01a04f1382dc26974b3c834052947c9 Mon Sep 17 00:00:00 2001 From: xiaoxiaohehe001 <49090790+xiaoxiaohehe001@users.noreply.github.com> Date: Fri, 17 Jun 2022 13:50:23 +0800 Subject: [PATCH] [Paddle Inference] Support dyanmic_shape of shuffle_channel. (#43488) * shuffle_channel * shuffle * shuffle * shuffle --- .../tensorrt/convert/shuffle_channel_op.cc | 65 ++++++++++++++----- paddle/fluid/inference/tensorrt/op_teller.cc | 5 +- .../test_trt_convert_shuffle_channel.py | 4 +- 3 files changed, 57 insertions(+), 17 deletions(-) diff --git a/paddle/fluid/inference/tensorrt/convert/shuffle_channel_op.cc b/paddle/fluid/inference/tensorrt/convert/shuffle_channel_op.cc index e6422522e50..ed113798a73 100644 --- a/paddle/fluid/inference/tensorrt/convert/shuffle_channel_op.cc +++ b/paddle/fluid/inference/tensorrt/convert/shuffle_channel_op.cc @@ -39,25 +39,60 @@ class ShuffleChannelOpConverter : public OpConverter { // Declare inputs auto* input = engine_->GetITensor(op_desc.Input("X")[0]); auto input_dims = input->getDimensions(); - - int c = input_dims.d[0]; - int h = input_dims.d[1]; - int w = input_dims.d[2]; + auto output_name = op_desc.Output("Out")[0]; int group = BOOST_GET_CONST(int, op_desc.GetAttr("group")); - auto* layer = TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *input); - nvinfer1::Dims4 reshape_dim(group, c / group, h, w); - layer->setReshapeDimensions(reshape_dim); - layer->setSecondTranspose({1, 0, 2, 3}); - auto* output = layer->getOutput(0); +#if IS_TRT_VERSION_GE(8000) + if (engine_->with_dynamic_shape()) { + auto* input_shape_tensor = Shape(input); + auto* channel_shape_tensor = GetEleTensorOfShape(input_shape_tensor, 1); + auto* group_tensor = + Add1DConstantLayer(group, output_name + "_group_tensor_"); + auto* new_channel_shape_tensor = Div(channel_shape_tensor, group_tensor); + std::vector shape_dim3{0, 2, 3}; + auto* shape_dim3_tensor = Gather(input_shape_tensor, shape_dim3); - auto* reshape_layer = TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *output); - nvinfer1::Dims3 reshape_dim2(c, h, w); - reshape_layer->setReshapeDimensions(reshape_dim2); + std::vector itensors; + itensors.push_back(shape_dim3_tensor); + itensors.push_back(group_tensor); + itensors.push_back(new_channel_shape_tensor); + auto* reshape_tensor = Concat(itensors); - auto output_name = op_desc.Output("Out")[0]; - RreplenishLayerAndOutput(reshape_layer, "shuffle_channel", {output_name}, - test_mode); + auto* reshape_layer = + TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *reshape_tensor); + nvinfer1::Permutation transpose_new_input{0, 3, 4, 1, 2}; + reshape_layer->setSecondTranspose(transpose_new_input); + + auto* layer = TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *input); + layer->setInput(1, *(reshape_layer->getOutput(0))); + nvinfer1::Permutation transpose_embed{0, 2, 1, 3, 4}; + layer->setSecondTranspose(transpose_embed); + auto* output = layer->getOutput(0); + auto* output_layer = TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *output); + output_layer->setInput(1, *input_shape_tensor); + + RreplenishLayerAndOutput(output_layer, "shuffle_channel", {output_name}, + test_mode); + } +#endif + if (!engine_->with_dynamic_shape()) { + int c = input_dims.d[0]; + int h = input_dims.d[1]; + int w = input_dims.d[2]; + + auto* layer = TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *input); + nvinfer1::Dims4 reshape_dim(group, c / group, h, w); + layer->setReshapeDimensions(reshape_dim); + layer->setSecondTranspose({1, 0, 2, 3}); + auto* output = layer->getOutput(0); + + auto* reshape_layer = TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *output); + nvinfer1::Dims3 reshape_dim2(c, h, w); + reshape_layer->setReshapeDimensions(reshape_dim2); + + RreplenishLayerAndOutput(reshape_layer, "shuffle_channel", {output_name}, + test_mode); + } } }; diff --git a/paddle/fluid/inference/tensorrt/op_teller.cc b/paddle/fluid/inference/tensorrt/op_teller.cc index d7e481bfa91..d6aa04612d6 100644 --- a/paddle/fluid/inference/tensorrt/op_teller.cc +++ b/paddle/fluid/inference/tensorrt/op_teller.cc @@ -1604,11 +1604,14 @@ bool OpTeller::Tell(const framework::ir::Node* node, bool use_no_calib_int8, } if (op_type == "shuffle_channel") { +#if !IS_TRT_VERSION_GE(8000) if (with_dynamic_shape) { VLOG(3) << "You are running the TRT Dynamic Shape mode, " - "the shuffle_channel op does not support dynamic shape yet"; + "the shuffle_channel op does not support dynamic shape " + "trt versions below 8.0 yet"; return false; } +#endif } if (op_type == "skip_layernorm") { diff --git a/python/paddle/fluid/tests/unittests/ir/inference/test_trt_convert_shuffle_channel.py b/python/paddle/fluid/tests/unittests/ir/inference/test_trt_convert_shuffle_channel.py index 9948b29321d..a53b61a0072 100644 --- a/python/paddle/fluid/tests/unittests/ir/inference/test_trt_convert_shuffle_channel.py +++ b/python/paddle/fluid/tests/unittests/ir/inference/test_trt_convert_shuffle_channel.py @@ -77,7 +77,9 @@ class TrtConvertShuffleChannelTest(TrtLayerAutoScanTest): self.dynamic_shape.opt_input_shape = {} def generate_trt_nodes_num(attrs, dynamic_shape): - if dynamic_shape == True: + ver = paddle_infer.get_trt_compile_version() + if ver[0] * 1000 + ver[1] * 100 + ver[ + 2] * 10 < 8000 and dynamic_shape == True: return 0, 3 else: return 1, 2 -- GitLab