[Paddle-TRT][cherry-pick]squeeze/unsqueeze -> 2.3

paddle/fluid/inference/api/analysis_predictor.cc

@@ -1808,6 +1808,8 @@ USE_TRT_CONVERTER(fused_preln_embedding_eltwise_layernorm)
 USE_TRT_CONVERTER(preln_skip_layernorm)
 USE_TRT_CONVERTER(roll)
 USE_TRT_CONVERTER(strided_slice)
+USE_TRT_CONVERTER(squeeze2)
+USE_TRT_CONVERTER(unsqueeze2)
 #endif
 
 namespace paddle_infer {

paddle/fluid/inference/tensorrt/convert/CMakeLists.txt

@@ -56,6 +56,8 @@ nv_library(
        strided_slice_op.cc
        preln_skip_layernorm.cc
        roll_op.cc
+       squeeze2_op.cc
+       unsqueeze2_op.cc
   DEPS tensorrt_engine tensorrt_plugin operator scope framework_proto
        op_registry)
 

paddle/fluid/inference/tensorrt/convert/squeeze2_op.cc

@@ -0,0 +1,82 @@
+/* Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#include "paddle/fluid/inference/tensorrt/convert/op_converter.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+
+class Squeeze2OpConverter : public OpConverter {
+ public:
+  void operator()(const framework::proto::OpDesc& op,
+                  const framework::Scope& scope,
+                  bool test_mode) override {
+    VLOG(4) << "convert a fluid squeeze2 op to tensorrt shuffle layer";
+
+    framework::OpDesc op_desc(op, nullptr);
+    // Declare inputs
+    auto* input = engine_->GetITensor(op_desc.Input("X")[0]);
+    auto input_dims = input->getDimensions();
+    auto output_name = op_desc.Output("Out")[0];
+
+    // Get Attrs
+    std::vector<int> axes =
+        BOOST_GET_CONST(std::vector<int>, op_desc.GetAttr("axes"));
+    PADDLE_ENFORCE_GT(
+        axes.size(),
+        0,
+        platform::errors::InvalidArgument(
+            "Attr(axes).size should be > 0 in squeeze2 op in TensorRT,"
+            "but received axes.size() = %d.",
+            axes.size()));
+
+    std::vector<bool> should_squeeze(input_dims.nbDims, false);
+    for (size_t i = 0; i < axes.size(); i++) {
+      if (engine_->with_dynamic_shape()) {
+        axes[i] += (axes[i] < 0) ? input_dims.nbDims : 0;
+      } else {
+        axes[i] += (axes[i] < 0) ? input_dims.nbDims : -1;
+      }
+      should_squeeze[axes[i]] = true;
+    }
+
+    nvinfer1::Dims trt_out_dims;
+    trt_out_dims.nbDims = 0;
+    std::vector<int32_t> gather_indices;
+    for (size_t i = 0; i < should_squeeze.size(); i++) {
+      if (should_squeeze[i]) continue;
+      gather_indices.push_back(i);
+      // for static shape
+      trt_out_dims.d[trt_out_dims.nbDims] = input_dims.d[i];
+      trt_out_dims.nbDims++;
+    }
+
+    auto* layer = TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *input);
+    if (engine_->with_dynamic_shape()) {
+      auto* shape_tensor = Shape(input);
+      auto* real_shape_tensor = Gather(shape_tensor, gather_indices);
+      layer->setInput(1, *real_shape_tensor);
+    } else {
+      layer->setReshapeDimensions(trt_out_dims);
+    }
+    RreplenishLayerAndOutput(layer, "squeeze2", {output_name}, test_mode);
+  }
+};
+
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle
+
+REGISTER_TRT_OP_CONVERTER(squeeze2, Squeeze2OpConverter);

paddle/fluid/inference/tensorrt/convert/unsqueeze2_op.cc

@@ -0,0 +1,101 @@
+/* Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#include "paddle/fluid/inference/tensorrt/convert/op_converter.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+
+class Unsqueeze2OpConverter : public OpConverter {
+ public:
+  void operator()(const framework::proto::OpDesc& op,
+                  const framework::Scope& scope,
+                  bool test_mode) override {
+    VLOG(4) << "convert a fluid unsqueeze2 op to tensorrt shuffle layer";
+
+    framework::OpDesc op_desc(op, nullptr);
+    // Declare inputs
+    auto* input = engine_->GetITensor(op_desc.Input("X")[0]);
+    auto input_dims = input->getDimensions();
+    auto output_name = op_desc.Output("Out")[0];
+
+    // Get Attrs
+    std::vector<int> axes =
+        BOOST_GET_CONST(std::vector<int>, op_desc.GetAttr("axes"));
+    PADDLE_ENFORCE_GT(
+        axes.size(),
+        0,
+        platform::errors::InvalidArgument(
+            "Attr(axes).size should be > 0 in unsqueeze2 op in TensorRT,"
+            "but received axes.size() = %d.",
+            axes.size()));
+
+    std::vector<bool> should_unsqueeze(input_dims.nbDims + axes.size(), false);
+    int cur_out_rank = input_dims.nbDims;
+    for (size_t i = 0; i < axes.size(); i++) {
+      cur_out_rank++;
+      if (engine_->with_dynamic_shape()) {
+        axes[i] += (axes[i] < 0) ? cur_out_rank : 0;
+      } else {
+        axes[i] += (axes[i] < 0) ? cur_out_rank : -1;
+      }
+      // axes[i] is relative to cur_out_rank
+      // we make [axes[i], cur_out_rank - 2] shift right
+      // and make (axes[i]) to true!
+      for (int j = cur_out_rank - 1; j > axes[i]; j--) {
+        should_unsqueeze[j] = should_unsqueeze[j - 1];
+      }
+      if (axes[i] >= cur_out_rank)
+        should_unsqueeze[cur_out_rank - 1] = true;
+      else
+        should_unsqueeze[axes[i]] = true;
+    }
+
+    nvinfer1::Dims trt_out_dims;
+    trt_out_dims.nbDims = should_unsqueeze.size();
+    std::vector<int32_t> gather_indices;
+    int in_rank_i = 0;
+    for (size_t i = 0; i < should_unsqueeze.size(); i++) {
+      if (should_unsqueeze[i]) {
+        trt_out_dims.d[i] = 1;
+        gather_indices.push_back(input_dims.nbDims);
+        continue;
+      }
+      trt_out_dims.d[i] = input_dims.d[in_rank_i];
+      gather_indices.push_back(in_rank_i);
+      in_rank_i++;
+    }
+
+    auto* layer = TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *input);
+    if (engine_->with_dynamic_shape()) {
+      auto* shape_tensor = Shape(input);
+      std::vector<int32_t> all_one(axes.size(), 1);
+      auto* all_one_tensor = Add1DConstantLayer(all_one);
+      std::vector<nvinfer1::ITensor*> concat_inputs = {shape_tensor,
+                                                       all_one_tensor};
+      auto* real_shape_tensor = Gather(Concat(concat_inputs), gather_indices);
+      layer->setInput(1, *real_shape_tensor);
+    } else {
+      layer->setReshapeDimensions(trt_out_dims);
+    }
+    RreplenishLayerAndOutput(layer, "unsqueeze2", {output_name}, test_mode);
+  }
+};
+
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle
+
+REGISTER_TRT_OP_CONVERTER(unsqueeze2, Unsqueeze2OpConverter);

paddle/fluid/inference/tensorrt/op_teller.cc

@@ -114,6 +114,8 @@ struct SimpleOpTypeSetTeller : public Teller {
       "bilinear_interp_v2",
       "cast",
       "pool3d",
+      "squeeze2",
+      "unsqueeze2",
       "deformable_conv",
       "relu6",
       "hard_sigmoid",
@@ -179,6 +181,8 @@ struct SimpleOpTypeSetTeller : public Teller {
       "nearest_interp_v2",
       "cast",
       "pool3d",
+      "squeeze2",
+      "unsqueeze2",
       "deformable_conv",
       "relu6",
       "hard_sigmoid",
@@ -891,6 +895,44 @@ bool OpTeller::Tell(const framework::ir::Node* node,
       }
     }
 
+    if (op_type == "squeeze2") {
+      std::vector<int> axes;
+      if (desc.HasAttr("axes")) {
+        axes = BOOST_GET_CONST(std::vector<int>, desc.GetAttr("axes"));
+      }
+      if (axes.size() == 0) {
+        VLOG(3) << "The necessary attributes of the squeeze2 operator axes is "
+                   "missing.";
+        return false;
+      }
+      if (!with_dynamic_shape) {
+        if (std::find(axes.begin(), axes.end(), 0) != axes.end()) {
+          VLOG(3) << "Invalid squeeze axes. Axes having batch axis is not "
+                     "supported in static shape";
+          return false;
+        }
+      }
+    }
+
+    if (op_type == "unsqueeze2") {
+      std::vector<int> axes;
+      if (desc.HasAttr("axes")) {
+        axes = BOOST_GET_CONST(std::vector<int>, desc.GetAttr("axes"));
+      }
+      if (axes.size() == 0) {
+        VLOG(3) << "The necessary attributes of the squeeze2 operator axes is "
+                   "missing.";
+        return false;
+      }
+      if (!with_dynamic_shape) {
+        if (std::find(axes.begin(), axes.end(), 0) != axes.end()) {
+          VLOG(3) << "Invalid squeeze axes. Axes having batch axis is not "
+                     "supported in static shape";
+          return false;
+        }
+      }
+    }
+
     if (op_type == "batch_norm") {
       const std::vector<std::string> bn_inputs = {
           "X", "Bias", "Mean", "Scale", "Variance"};