diff --git a/tools/infer/utility.py b/tools/infer/utility.py index f4e62b87a506ccff889df398d363411a919b5ed5..8ad916bee8ec6872a78448446447b6b7f6086a56 100755 --- a/tools/infer/utility.py +++ b/tools/infer/utility.py @@ -139,9 +139,83 @@ def create_predictor(args, mode, logger): config.enable_use_gpu(args.gpu_mem, 0) if args.use_tensorrt: config.enable_tensorrt_engine( - precision_mode=inference.PrecisionType.Half - if args.use_fp16 else inference.PrecisionType.Float32, - max_batch_size=args.max_batch_size) + precision_mode=inference.PrecisionType.Float32, + max_batch_size=args.max_batch_size, + min_subgraph_size=3) # skip the minmum trt subgraph + if mode == "det" and "mobile" in model_file_path: + min_input_shape = { + "x": [1, 3, 50, 50], + "conv2d_92.tmp_0": [1, 96, 20, 20], + "conv2d_91.tmp_0": [1, 96, 10, 10], + "nearest_interp_v2_1.tmp_0": [1, 96, 10, 10], + "nearest_interp_v2_2.tmp_0": [1, 96, 20, 20], + "nearest_interp_v2_3.tmp_0": [1, 24, 20, 20], + "nearest_interp_v2_4.tmp_0": [1, 24, 20, 20], + "nearest_interp_v2_5.tmp_0": [1, 24, 20, 20], + "elementwise_add_7": [1, 56, 2, 2], + "nearest_interp_v2_0.tmp_0": [1, 96, 2, 2] + } + max_input_shape = { + "x": [1, 3, 2000, 2000], + "conv2d_92.tmp_0": [1, 96, 400, 400], + "conv2d_91.tmp_0": [1, 96, 200, 200], + "nearest_interp_v2_1.tmp_0": [1, 96, 200, 200], + "nearest_interp_v2_2.tmp_0": [1, 96, 400, 400], + "nearest_interp_v2_3.tmp_0": [1, 24, 400, 400], + "nearest_interp_v2_4.tmp_0": [1, 24, 400, 400], + "nearest_interp_v2_5.tmp_0": [1, 24, 400, 400], + "elementwise_add_7": [1, 56, 400, 400], + "nearest_interp_v2_0.tmp_0": [1, 96, 400, 400] + } + opt_input_shape = { + "x": [1, 3, 640, 640], + "conv2d_92.tmp_0": [1, 96, 160, 160], + "conv2d_91.tmp_0": [1, 96, 80, 80], + "nearest_interp_v2_1.tmp_0": [1, 96, 80, 80], + "nearest_interp_v2_2.tmp_0": [1, 96, 160, 160], + "nearest_interp_v2_3.tmp_0": [1, 24, 160, 160], + "nearest_interp_v2_4.tmp_0": [1, 24, 160, 160], + "nearest_interp_v2_5.tmp_0": [1, 24, 160, 160], + "elementwise_add_7": [1, 56, 40, 40], + "nearest_interp_v2_0.tmp_0": [1, 96, 40, 40] + } + if mode == "det" and "server" in model_file_path: + min_input_shape = { + "x": [1, 3, 50, 50], + "conv2d_59.tmp_0": [1, 96, 20, 20], + "nearest_interp_v2_2.tmp_0": [1, 96, 20, 20], + "nearest_interp_v2_3.tmp_0": [1, 24, 20, 20], + "nearest_interp_v2_4.tmp_0": [1, 24, 20, 20], + "nearest_interp_v2_5.tmp_0": [1, 24, 20, 20] + } + max_input_shape = { + "x": [1, 3, 2000, 2000], + "conv2d_59.tmp_0": [1, 96, 400, 400], + "nearest_interp_v2_2.tmp_0": [1, 96, 400, 400], + "nearest_interp_v2_3.tmp_0": [1, 24, 400, 400], + "nearest_interp_v2_4.tmp_0": [1, 24, 400, 400], + "nearest_interp_v2_5.tmp_0": [1, 24, 400, 400] + } + opt_input_shape = { + "x": [1, 3, 640, 640], + "conv2d_59.tmp_0": [1, 96, 160, 160], + "nearest_interp_v2_2.tmp_0": [1, 96, 160, 160], + "nearest_interp_v2_3.tmp_0": [1, 24, 160, 160], + "nearest_interp_v2_4.tmp_0": [1, 24, 160, 160], + "nearest_interp_v2_5.tmp_0": [1, 24, 160, 160] + } + elif mode == "rec": + min_input_shape = {"x": [args.rec_batch_num, 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]} + max_input_shape = {"x": [args.rec_batch_num, 3, 48, 2000]} + opt_input_shape = {"x": [args.rec_batch_num, 3, 48, 320]} + + config.set_trt_dynamic_shape_info(min_input_shape, max_input_shape, + opt_input_shape) + else: config.disable_gpu() if hasattr(args, "cpu_threads"):