ONNX 模型文件转换为 Python脚本

ONNX 模型文件转换为 Python脚本#

onnx2script 模块提供一个脚本,用于使用 onnxscript 库将 ONNX 模型文件转换为 Python 脚本。

import os
import onnx
import onnxscript

def convert2script(
    input_file_name: str, output_file_name: str|None=None, verbose: bool=True,
    load_external_data=False,
) -> None:
    """使用 `onnxscript` 库将 ONNX 模型文件转换为 Python 脚本"""
    model = onnx.load(input_file_name, load_external_data=load_external_data)
    python_code = onnxscript.proto2python(
        model, use_operators=not verbose, inline_const=not verbose
    )

    # 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 + ".py"

    with open(output_file_name, "w", encoding="utf-8") as f:
        f.write(python_code)
input_file_name = "/media/pc/data/board/arria10/lxw/tasks/tools/npuusertools/models/telecom/onnx/Nin1_helmet_small/helmet_small.onnx"
convert2script(input_file_name)