codekingpro/portable-devtools

1# --------------------------------------------------------------------------
2# Copyright (c) Microsoft, Intel Corporation. All rights reserved.
3# Licensed under the MIT License. See License.txt in the project root for
4# license information.
5# --------------------------------------------------------------------------
6
7import argparse
8import logging
9import sys
10
11from .shape_inference import quant_pre_process
12
13logger = logging.getLogger(__name__)
14
15
16def parse_arguments():
17    parser = argparse.ArgumentParser(
18        description="""Model optimizer and shape inferencer, in preparation for quantization,
19Consists of three optional steps:
201. Symbolic shape inference (best for transformer models).
212. Model optimization.
223. ONNX shape inference.
23
24Model quantization with QDQ format, i.e. inserting QuantizeLinear/DeQuantizeLinear on
25the tensor, requires tensor shape information to perform its best. Currently, shape inferencing
26works best with optimized model. As a result, it is highly recommended to run quantization
27on optimized model with shape information. This is the tool for optimization and shape
28inferencing.
29
30Essentially this tool performs the following three (skippable) steps:
31
321. Symbolic shape inference.
332. Model optimization
343. ONNX shape inference"""
35    )
36
37    parser.add_argument("--input", required=True, help="Path to the input model file")
38    parser.add_argument("--output", required=True, help="Path to the output model file")
39    parser.add_argument(
40        "--skip_optimization",
41        type=bool,
42        default=False,
43        help="Skip model optimization step if true. It's a known issue that ORT"
44        " optimization has difficulty with model size greater than 2GB, rerun with"
45        " this option to get around this issue.",
46    )
47    parser.add_argument(
48        "--skip_onnx_shape",
49        type=bool,
50        default=False,
51        help="Skip ONNX shape inference. Symbolic shape inference is most effective"
52        " with transformer based models. Skipping all shape inferences may"
53        " reduce the effectiveness of quantization, as a tensor with unknown"
54        " shape can not be quantized.",
55    )
56    parser.add_argument(
57        "--skip_symbolic_shape",
58        type=bool,
59        default=False,
60        help="Skip symbolic shape inference. Symbolic shape inference is most"
61        " effective with transformer based models. Skipping all shape"
62        " inferences may reduce the effectiveness of quantization, as a tensor"
63        " with unknown shape can not be quantized.",
64    )
65    parser.add_argument(
66        "--auto_merge",
67        help="Automatically merge symbolic dims when confliction happens",
68        action="store_true",
69        default=False,
70    )
71    parser.add_argument(
72        "--int_max",
73        help="maximum value for integer to be treated as boundless for ops like slice",
74        type=int,
75        default=2**31 - 1,
76    )
77    parser.add_argument(
78        "--guess_output_rank",
79        help="guess output rank to be the same as input 0 for unknown ops",
80        action="store_true",
81        default=False,
82    )
83    parser.add_argument(
84        "--verbose",
85        help="Prints detailed logs of inference, 0: turn off, 1: warnings, 3: detailed",
86        type=int,
87        default=0,
88    )
89    parser.add_argument(
90        "--save_as_external_data",
91        help="Saving an ONNX model to external data",
92        action="store_true",
93        default=False,
94    )
95    parser.add_argument(
96        "--all_tensors_to_one_file",
97        help="Saving all the external data to one file",
98        action="store_true",
99        default=False,
100    )
101    parser.add_argument(
102        "--external_data_location",
103        help="The file location to save the external file",
104        default=None,
105    )
106    parser.add_argument(
107        "--external_data_size_threshold",
108        help="The size threshold for external data",
109        type=int,
110        default=1024,
111    )
112    return parser.parse_args()
113
114
115if __name__ == "__main__":
116    args = parse_arguments()
117    if args.skip_optimization and args.skip_onnx_shape and args.skip_symbolic_shape:
118        logger.error("Skipping all three steps, nothing to be done. Quitting...")
119        sys.exit()
120
121    if (not args.skip_optimization) and args.save_as_external_data:
122        logger.error("ORT model optimization does not support external data yet!")
123        sys.exit()
124
125    logger.info("input model: %s", args.input)
126    logger.info("output model: %s", args.output)
127    quant_pre_process(
128        args.input,
129        args.output,
130        args.skip_optimization,
131        args.skip_onnx_shape,
132        args.skip_symbolic_shape,
133        args.auto_merge,
134        args.int_max,
135        args.guess_output_rank,
136        args.verbose,
137        args.save_as_external_data,
138        args.all_tensors_to_one_file,
139        args.external_data_location,
140        args.external_data_size_threshold,
141    )
142