mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-05-15 20:50:42 +00:00
02c7873b0e
* Add support for custom ops library to the ORT model conversion script Simplify model conversion now that we read ops from the ORT format model. Enable custom ops in the python bindings if custom ops are turned on in a minimal build. * Add test of model conversion involving custom ops.
78 lines
3.5 KiB
Python
78 lines
3.5 KiB
Python
# Copyright (c) Microsoft Corporation. All rights reserved.
|
|
# Licensed under the MIT License.
|
|
|
|
import os
|
|
|
|
from .operator_type_usage_processors import OperatorTypeUsageManager
|
|
from .ort_model_processor import OrtFormatModelProcessor
|
|
|
|
from ..logger import get_logger
|
|
log = get_logger("ort_format_model.utils")
|
|
|
|
|
|
def _extract_ops_and_types_from_ort_models(model_path_or_dir: str, enable_type_reduction: bool):
|
|
if not os.path.exists(model_path_or_dir):
|
|
raise ValueError('Path to model/s does not exist: {}'.format(model_path_or_dir))
|
|
|
|
required_ops = {}
|
|
op_type_usage_manager = OperatorTypeUsageManager() if enable_type_reduction else None
|
|
|
|
if os.path.isfile(model_path_or_dir):
|
|
model_processor = OrtFormatModelProcessor(model_path_or_dir, required_ops, op_type_usage_manager)
|
|
model_processor.process() # this updates required_ops and op_type_processors
|
|
log.info('Processed {}'.format(model_path_or_dir))
|
|
else:
|
|
for root, _, files in os.walk(model_path_or_dir):
|
|
for file in files:
|
|
model_path = os.path.join(root, file)
|
|
if file.lower().endswith('.ort'):
|
|
model_processor = OrtFormatModelProcessor(model_path, required_ops, op_type_usage_manager)
|
|
model_processor.process() # this updates required_ops and op_type_processors
|
|
log.info('Processed {}'.format(model_path))
|
|
|
|
return required_ops, op_type_usage_manager
|
|
|
|
|
|
def create_config_from_models(model_path_or_dir: str, output_file: str = None, enable_type_reduction: bool = True):
|
|
'''
|
|
Create a configuration file with required operators and optionally required types.
|
|
:param model_path_or_dir: Path to recursively search for ORT format models, or to a single ORT format model.
|
|
:param output_file: File to write configuration to.
|
|
Defaults to creating required_operators.config in the model_path_or_dir directory.
|
|
:param enable_type_reduction: Include required type information for individual operators in the configuration.
|
|
'''
|
|
|
|
required_ops, op_type_processors = _extract_ops_and_types_from_ort_models(model_path_or_dir, enable_type_reduction)
|
|
|
|
if output_file:
|
|
directory, filename = os.path.split(output_file)
|
|
if not filename:
|
|
raise RuntimeError("Invalid output path for configuration: {}".format(output_file))
|
|
|
|
if directory and not os.path.exists(directory):
|
|
os.makedirs(directory)
|
|
else:
|
|
dir = model_path_or_dir
|
|
if os.path.isfile(model_path_or_dir):
|
|
dir = os.path.basename(model_path_or_dir)
|
|
|
|
output_file = os.path.join(dir, 'required_operators.config')
|
|
|
|
with open(output_file, 'w') as out:
|
|
out.write("# Generated from model/s in {}\n".format(model_path_or_dir))
|
|
|
|
for domain in sorted(required_ops.keys()):
|
|
for opset in sorted(required_ops[domain].keys()):
|
|
ops = required_ops[domain][opset]
|
|
if ops:
|
|
out.write("{};{};".format(domain, opset))
|
|
if enable_type_reduction:
|
|
# type string is empty if op hasn't been seen
|
|
entries = ['{}{}'.format(op, op_type_processors.get_config_entry(domain, op) or '')
|
|
for op in sorted(ops)]
|
|
else:
|
|
entries = sorted(ops)
|
|
|
|
out.write("{}\n".format(','.join(entries)))
|
|
|
|
log.info("Created config in %s", output_file)
|