提交 8f838187 编写于 作者: L LDOUBLEV

fix trt

上级 d4c4877e
......@@ -195,7 +195,8 @@ def create_predictor(args, mode, logger):
max_batch_size=args.max_batch_size,
min_subgraph_size=args.min_subgraph_size)
# skip the minmum trt subgraph
if mode == "det":
use_dynamic_shape = True
if mode == "det" and args.rec_algorithm == "DB":
min_input_shape = {
"x": [1, 3, 50, 50],
"conv2d_92.tmp_0": [1, 120, 20, 20],
......@@ -260,6 +261,8 @@ def create_predictor(args, mode, logger):
max_input_shape.update(max_pact_shape)
opt_input_shape.update(opt_pact_shape)
elif mode == "rec":
if args.rec_algorithm != "CRNN":
use_dynamic_shape = False
min_input_shape = {"x": [1, 3, 32, 10]}
max_input_shape = {"x": [args.rec_batch_num, 3, 32, 1536]}
opt_input_shape = {"x": [args.rec_batch_num, 3, 32, 320]}
......@@ -268,14 +271,8 @@ def create_predictor(args, mode, logger):
max_input_shape = {"x": [args.rec_batch_num, 3, 48, 1024]}
opt_input_shape = {"x": [args.rec_batch_num, 3, 48, 320]}
else:
min_input_shape = {"x": [1, 3, 10, 10]}
max_input_shape = {"x": [1, 3, 512, 512]}
opt_input_shape = {"x": [1, 3, 256, 256]}
if mode == "rec":
if args.rec_algorithm == "CRNN":
config.set_trt_dynamic_shape_info(
min_input_shape, max_input_shape, opt_input_shape)
else:
use_dynamic_shape = False
if use_dynamic_shape:
config.set_trt_dynamic_shape_info(
min_input_shape, max_input_shape, opt_input_shape)
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册