From 96a53fb4becfface113ec4e717b2cfa976764c63 Mon Sep 17 00:00:00 2001 From: LDOUBLEV Date: Thu, 10 Jun 2021 14:29:06 +0800 Subject: [PATCH] add params min_subgraph_size --- tools/infer/utility.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tools/infer/utility.py b/tools/infer/utility.py index 2b70f36e..a87e3eef 100755 --- a/tools/infer/utility.py +++ b/tools/infer/utility.py @@ -37,6 +37,7 @@ def init_args(): parser.add_argument("--use_gpu", type=str2bool, default=True) parser.add_argument("--ir_optim", type=str2bool, default=True) parser.add_argument("--use_tensorrt", type=str2bool, default=False) + parser.add_argument("--min_subgraph_size", type=int, default=3) parser.add_argument("--precision", type=str, default="fp32") parser.add_argument("--gpu_mem", type=int, default=500) @@ -235,7 +236,8 @@ def create_predictor(args, mode, logger): config.enable_tensorrt_engine( precision_mode=inference.PrecisionType.Float32, max_batch_size=args.max_batch_size, - min_subgraph_size=3) # skip the minmum trt subgraph + min_subgraph_size=args. + min_subgraph_size) # skip the minmum trt subgraph if mode == "det": min_input_shape = { "x": [1, 3, 50, 50],