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

fix trt

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