From e79667ac885eb93444cf35c691c4d249d52f0f50 Mon Sep 17 00:00:00 2001 From: LDOUBLEV Date: Wed, 8 Sep 2021 10:24:30 +0800 Subject: [PATCH] support dynamic batch shape for trt prediction --- tools/infer/utility.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tools/infer/utility.py b/tools/infer/utility.py index 7f60773c..466f824c 100755 --- a/tools/infer/utility.py +++ b/tools/infer/utility.py @@ -236,11 +236,11 @@ def create_predictor(args, mode, logger): max_input_shape.update(max_pact_shape) opt_input_shape.update(opt_pact_shape) elif mode == "rec": - min_input_shape = {"x": [args.rec_batch_num, 3, 32, 10]} + min_input_shape = {"x": [1, 3, 32, 10]} max_input_shape = {"x": [args.rec_batch_num, 3, 32, 2000]} opt_input_shape = {"x": [args.rec_batch_num, 3, 32, 320]} elif mode == "cls": - min_input_shape = {"x": [args.rec_batch_num, 3, 48, 10]} + min_input_shape = {"x": [1, 3, 48, 10]} max_input_shape = {"x": [args.rec_batch_num, 3, 48, 2000]} opt_input_shape = {"x": [args.rec_batch_num, 3, 48, 320]} else: -- GitLab