[Paddle Inference] Support dyanmic_shape of shuffle_channel. (#43488)

* shuffle_channel

* shuffle

* shuffle

* shuffle

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<int32_t> 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<nvinfer1::ITensor*> 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);
+    }
   }
 };
 

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") {

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