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)