mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-05-15 20:50:42 +00:00
89 lines
4.1 KiB
Python
89 lines
4.1 KiB
Python
#!/usr/bin/env python
|
|
# Copyright (c) Microsoft Corporation. All rights reserved.
|
|
# Licensed under the MIT License.
|
|
|
|
import argparse
|
|
import glob
|
|
import os
|
|
import re
|
|
import sys
|
|
import tempfile
|
|
|
|
import onnxruntime as ort
|
|
|
|
|
|
def create_config_file(optimized_model_path, config_file_path):
|
|
script_path = os.path.dirname(os.path.realpath(__file__))
|
|
ci_build_py_path = os.path.abspath(os.path.join(script_path, '..', 'ci_build'))
|
|
sys.path.append(ci_build_py_path)
|
|
|
|
# create config file from all the optimized models
|
|
print("Creating configuration file for operators required by optimized models in {}".format(config_file_path))
|
|
from exclude_unused_ops import exclude_unused_ops # tools/ci_build/exclude_unused_ops.py
|
|
exclude_unused_ops(optimized_model_path, config_path=None, ort_root=None, output_config_path=config_file_path)
|
|
|
|
|
|
def convert(model_path: str):
|
|
models = glob.glob(os.path.join(model_path, '**', '*.onnx'), recursive=True)
|
|
|
|
if len(models) == 0:
|
|
raise ValueError("No .onnx files were found in " + model_path)
|
|
|
|
# create temp directory to create optimized onnx format models in. currently we need this to create the
|
|
# config file with required operators. long term we could potentially do this from the ORT format model,
|
|
# however that requires a lot of infrastructure to be able to parse the flatbuffers schema for those files
|
|
with tempfile.TemporaryDirectory() as tmpdirname:
|
|
for model in models:
|
|
model_filename = os.path.basename(model)
|
|
# create .optimized.onnx file in temp dir
|
|
onnx_target_path = os.path.join(tmpdirname, re.sub('.onnx$', '.optimized.onnx', model_filename))
|
|
# create .ort file in same dir as original onnx model
|
|
ort_target_path = re.sub('.onnx$', '.ort', model)
|
|
|
|
so = ort.SessionOptions()
|
|
so.optimized_model_filepath = onnx_target_path
|
|
so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_EXTENDED # Skip NCHWc optimizations
|
|
|
|
print("Optimizing ONNX model {}".format(model))
|
|
# creating the session will result in the optimized model being saved
|
|
_ = ort.InferenceSession(model, sess_options=so)
|
|
|
|
# Second, convert optimized ONNX model to ORT format
|
|
so.optimized_model_filepath = ort_target_path
|
|
so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL # Convert model as-is so we don't change the kernels in this step # noqa
|
|
so.add_session_config_entry('session.save_model_format', 'ORT')
|
|
|
|
print("Converting optimized ONNX model to ORT format model {}".format(ort_target_path))
|
|
_ = ort.InferenceSession(onnx_target_path, sess_options=so)
|
|
|
|
# orig_size = os.path.getsize(onnx_target_path)
|
|
# new_size = os.path.getsize(ort_target_path)
|
|
# print("Serialized {} to {}. Sizes: orig={} new={} diff={} new:old={:.4f}:1.0".format(
|
|
# onnx_target_path, ort_target_path, orig_size, new_size, new_size - orig_size, new_size / orig_size))
|
|
|
|
# now that all models are converted create the config file before the temp dir is deleted
|
|
create_config_file(tmpdirname, os.path.join(model_path, 'required_operators.config'))
|
|
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser(
|
|
os.path.basename(__file__),
|
|
description='''Convert the ONNX format model/s in the provided directory to ORT format models.
|
|
All files with a `.onnx` extension will be processed. For each one, an ORT format model will be created in the
|
|
same directory. A configuration file will also be created called `required_operators.config`, and will contain
|
|
the list of required operators for all converted models.
|
|
This configuration file should be used as input to the minimal build'''
|
|
)
|
|
|
|
parser.add_argument('model_path', help='Provide path to directory containing ONNX model/s to convert. '
|
|
'Files with .onnx extension will be processed.')
|
|
return parser.parse_args()
|
|
|
|
|
|
def main():
|
|
args = parse_args()
|
|
convert(args.model_path)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|