提交 4a15797e 编写于 作者: Z zlsh80826

revise trt ernie test to conver the stack op CI testi, test=develop

上级 b4ade548
...@@ -85,6 +85,8 @@ struct SimpleOpTypeSetTeller : public Teller { ...@@ -85,6 +85,8 @@ struct SimpleOpTypeSetTeller : public Teller {
"gelu", "gelu",
"layer_norm", "layer_norm",
"scale", "scale",
"slice",
"stack",
}; };
}; };
......
...@@ -90,7 +90,6 @@ void trt_ernie(bool with_fp16, std::vector<float> result) { ...@@ -90,7 +90,6 @@ void trt_ernie(bool with_fp16, std::vector<float> result) {
config.SwitchUseFeedFetchOps(false); config.SwitchUseFeedFetchOps(false);
int head_number = 12;
int batch = 1; int batch = 1;
int min_seq_len = 1; int min_seq_len = 1;
int max_seq_len = 128; int max_seq_len = 128;
...@@ -104,23 +103,23 @@ void trt_ernie(bool with_fp16, std::vector<float> result) { ...@@ -104,23 +103,23 @@ void trt_ernie(bool with_fp16, std::vector<float> result) {
{"read_file_0.tmp_0", min_shape}, {"read_file_0.tmp_0", min_shape},
{"read_file_0.tmp_1", min_shape}, {"read_file_0.tmp_1", min_shape},
{"read_file_0.tmp_2", 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<std::string, std::vector<int>> max_input_shape = { std::map<std::string, std::vector<int>> max_input_shape = {
{"read_file_0.tmp_0", max_shape}, {"read_file_0.tmp_0", max_shape},
{"read_file_0.tmp_1", max_shape}, {"read_file_0.tmp_1", max_shape},
{"read_file_0.tmp_2", 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<std::string, std::vector<int>> opt_input_shape = { std::map<std::string, std::vector<int>> opt_input_shape = {
{"read_file_0.tmp_0", opt_shape}, {"read_file_0.tmp_0", opt_shape},
{"read_file_0.tmp_1", opt_shape}, {"read_file_0.tmp_1", opt_shape},
{"read_file_0.tmp_2", 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; auto precision = AnalysisConfig::Precision::kFloat32;
if (with_fp16) { if (with_fp16) {
precision = AnalysisConfig::Precision::kHalf; 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, config.SetTRTDynamicShapeInfo(min_input_shape, max_input_shape,
opt_input_shape); opt_input_shape);
std::vector<float> out_data; std::vector<float> out_data;
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册