onnxruntime

1
#!/usr/bin/env python3
2
# Copyright (c) Microsoft Corporation. All rights reserved.
3
# Licensed under the MIT License.
4

5
import argparse
6
import os
7
import pathlib
8

9
from .onnx_model_utils import get_optimization_level, optimize_model
10

11

12
def optimize_model_helper():
13
    parser = argparse.ArgumentParser(
14
        f"{os.path.basename(__file__)}:{optimize_model_helper.__name__}",
15
        description="""
16
                                     Optimize an ONNX model using ONNX Runtime to the specified level.
17
                                     See https://onnxruntime.ai/docs/performance/model-optimizations/graph-optimizations.html for more
18
                                     details of the optimization levels.""",
19
    )
20

21
    parser.add_argument(
22
        "--opt_level",
23
        default="basic",
24
        choices=["disable", "basic", "extended", "all"],
25
        help="Optimization level to use.",
26
    )
27
    parser.add_argument(
28
        "--log_level",
29
        choices=["debug", "info", "warning", "error"],
30
        type=str,
31
        required=False,
32
        default="error",
33
        help="Log level. Defaults to Error so we don't get output about unused initializers "
34
        "being removed. Warning or Info may be desirable in some scenarios.",
35
    )
36

37
    parser.add_argument("input_model", type=pathlib.Path, help="Provide path to ONNX model to update.")
38
    parser.add_argument("output_model", type=pathlib.Path, help="Provide path to write optimized ONNX model to.")
39

40
    args = parser.parse_args()
41

42
    if args.log_level == "error":
43
        log_level = 3
44
    elif args.log_level == "debug":
45
        log_level = 0  # ORT verbose level
46
    elif args.log_level == "info":
47
        log_level = 1
48
    elif args.log_level == "warning":
49
        log_level = 2
50

51
    optimize_model(args.input_model, args.output_model, get_optimization_level(args.opt_level), log_level)
52

53

54
if __name__ == "__main__":
55
    optimize_model_helper()
56