codekingpro/portable-devtools

1#!/usr/bin/env python3
2# Copyright (c) Microsoft Corporation. All rights reserved.
3# Licensed under the MIT License.
4from __future__ import annotations
5
6import argparse
7import os
8import pathlib
9
10from .onnx_model_utils import get_optimization_level, optimize_model
11
12
13def optimize_model_helper():
14    parser = argparse.ArgumentParser(
15        f"{os.path.basename(__file__)}:{optimize_model_helper.__name__}",
16        description="""
17                                     Optimize an ONNX model using ONNX Runtime to the specified level.
18                                     See https://onnxruntime.ai/docs/performance/model-optimizations/graph-optimizations.html for more
19                                     details of the optimization levels.""",
20    )
21
22    parser.add_argument(
23        "--opt_level",
24        default="basic",
25        choices=["disable", "basic", "extended", "layout", "all"],
26        help="Optimization level to use.",
27    )
28    parser.add_argument(
29        "--log_level",
30        choices=["debug", "info", "warning", "error"],
31        type=str,
32        required=False,
33        default="error",
34        help="Log level. Defaults to Error so we don't get output about unused initializers "
35        "being removed. Warning or Info may be desirable in some scenarios.",
36    )
37
38    parser.add_argument("input_model", type=pathlib.Path, help="Provide path to ONNX model to update.")
39    parser.add_argument("output_model", type=pathlib.Path, help="Provide path to write optimized ONNX model to.")
40
41    args = parser.parse_args()
42
43    if args.log_level == "error":
44        log_level = 3
45    elif args.log_level == "debug":
46        log_level = 0  # ORT verbose level
47    elif args.log_level == "info":
48        log_level = 1
49    elif args.log_level == "warning":
50        log_level = 2
51
52    optimize_model(args.input_model, args.output_model, get_optimization_level(args.opt_level), log_level)
53
54
55if __name__ == "__main__":
56    optimize_model_helper()
57