优化 ONNX 模型的实用工具

优化 ONNX 模型的实用工具#

参考:optimize.py

import os
import onnx
import onnx.inliner
import onnxscript

def optimize(
    input_file_name: str, output_file_name: str|None=None,
) -> None:
    """使用 `onnxscript` 库将 ONNX 模型文件转换为 Python 脚本"""
    model = onnx.load(input_file_name, load_external_data=False)
    model = onnxscript.optimizer.optimize(model)
    model = onnx.inliner.inline_local_functions(model)
    # Optimize again in case inlining created new opportunities.
    model = onnxscript.optimizer.optimize(model)

    # If output file name is not provided, use the input file name with .py extension
    if output_file_name is None:
        base_name = os.path.splitext(input_file_name)[0]  # Remove extension
        output_file_name = base_name + "_opt.onnx"

    onnx.save(model, output_file_name)
input_file_name = "/media/pc/data/board/arria10/lxw/tasks/tools/npuusertools/models/xmdemo/adas/_opt.onnx"
optimize(input_file_name)