diff --git a/paddle/fluid/inference/tensorrt/op_teller.cc b/paddle/fluid/inference/tensorrt/op_teller.cc index 70ead9720d2ebcb15ae0173dc0ba7c2095a4f4d4..3b7810d363fb6a5f3dc766b373b89ab713e6cd52 100644 --- a/paddle/fluid/inference/tensorrt/op_teller.cc +++ b/paddle/fluid/inference/tensorrt/op_teller.cc @@ -85,6 +85,8 @@ struct SimpleOpTypeSetTeller : public Teller { "gelu", "layer_norm", "scale", + "slice", + "stack", }; }; diff --git a/paddle/fluid/inference/tests/api/trt_dynamic_shape_ernie_test.cc b/paddle/fluid/inference/tests/api/trt_dynamic_shape_ernie_test.cc index babe9977cd571f588f0bdc5a6723d4b05afab72b..bf5f8828e0b65ccf0601b44e91094ea8e2d0dbfc 100644 --- a/paddle/fluid/inference/tests/api/trt_dynamic_shape_ernie_test.cc +++ b/paddle/fluid/inference/tests/api/trt_dynamic_shape_ernie_test.cc @@ -90,7 +90,6 @@ void trt_ernie(bool with_fp16, std::vector result) { config.SwitchUseFeedFetchOps(false); - int head_number = 12; int batch = 1; int min_seq_len = 1; int max_seq_len = 128; @@ -104,23 +103,23 @@ void trt_ernie(bool with_fp16, std::vector result) { {"read_file_0.tmp_0", min_shape}, {"read_file_0.tmp_1", min_shape}, {"read_file_0.tmp_2", min_shape}, - {"stack_0.tmp_0", {batch, head_number, min_seq_len, min_seq_len}}}; + {"matmul_0.tmp_0", min_shape}}; std::map> max_input_shape = { {"read_file_0.tmp_0", max_shape}, {"read_file_0.tmp_1", max_shape}, {"read_file_0.tmp_2", max_shape}, - {"stack_0.tmp_0", {batch, head_number, max_seq_len, max_seq_len}}}; + {"matmul_0.tmp_0", max_shape}}; std::map> opt_input_shape = { {"read_file_0.tmp_0", opt_shape}, {"read_file_0.tmp_1", opt_shape}, {"read_file_0.tmp_2", opt_shape}, - {"stack_0.tmp_0", {batch, head_number, opt_seq_len, opt_seq_len}}}; + {"matmul_0.tmp_0", opt_shape}}; auto precision = AnalysisConfig::Precision::kFloat32; if (with_fp16) { precision = AnalysisConfig::Precision::kHalf; } - config.EnableTensorRtEngine(1 << 30, 1, 5, precision, false, false); + config.EnableTensorRtEngine(1 << 30, 1, 3, precision, false, false); config.SetTRTDynamicShapeInfo(min_input_shape, max_input_shape, opt_input_shape); std::vector out_data;