reshape_opteller (#41090)

fix_reshape: for paddle-trt
PaddlePaddle · Apr 1, 2022 · 15d5f6b · 15d5f6b
1 parent f3270fc
commit 15d5f6b
Showing 1 changed file with 20 additions and 1 deletion.
diff --git a/paddle/fluid/inference/tensorrt/op_teller.cc b/paddle/fluid/inference/tensorrt/op_teller.cc
@@ -1479,8 +1479,27 @@ bool OpTeller::Tell(const framework::ir::Node* node, bool use_no_calib_int8,
       std::vector<int> shape =
           BOOST_GET_CONST(std::vector<int>, desc.GetAttr("shape"));
       if (shape.size() >= nvinfer1::Dims::MAX_DIMS) return false;
-      if (!with_dynamic_shape && (shape[0] == -1 || shape.size() == 1))
+      if (!with_dynamic_shape) {
+        if (shape.size() == 1) {
+          return false;
+        }
+        if (shape[0] == 0) {
+          return true;
+        } else {
+          auto* block = desc.Block();
+          auto x_var_name = desc.Input("X")[0];
+          auto* x_var_desc = block->FindVar(x_var_name);
+          const auto x_shape = x_var_desc->GetShape();
+          int input_num = std::accumulate(x_shape.begin() + 1, x_shape.end(), 1,
+                                          std::multiplies<int>());
+          int shape_num = std::accumulate(shape.begin() + 1, shape.end(), 1,
+                                          std::multiplies<int>());
+          if (input_num == shape_num) {
+            return true;
+          }
+        }
         return false;
+      }
     }
 
     if (op_type == "clip") {