mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-29 20:14:01 +00:00
Add ability to generate configuration file with required operators. (#5089)
* Add ability to generate configuration file with required operators.
This commit is contained in:
parent
80ada0291f
commit
dbf4e7019d
3 changed files with 110 additions and 47 deletions
|
|
@ -69,3 +69,9 @@ python exclude_unused_ops.py --model_path d:\ReduceSize\models --config_path d:\
|
|||
|
||||
After running the script build ORT as per the build instructions. Remember to specify `--skip-tests`.
|
||||
|
||||
#### Generating configuration file
|
||||
|
||||
Note: It is also possible to generate a configuration file for future usage by providing the `--write_combined_config_to` argument to `exclude_unused_ops.py`.
|
||||
If run this way it will process the information provided by `--model_path` and/or `--config_path`, and output a configuration file with the combined list of required operators to the provided path.
|
||||
No kernel registration changes will be made when run this way.
|
||||
|
||||
|
|
|
|||
|
|
@ -1631,10 +1631,10 @@ def main():
|
|||
args.test = False
|
||||
|
||||
if args.include_ops_by_model or args.include_ops_by_config:
|
||||
from exclude_unused_ops import exclude_unused_ops, get_provider_path
|
||||
include_ops_by_model = args.include_ops_by_model if args.include_ops_by_model else ''
|
||||
include_ops_by_config = args.include_ops_by_config if args.include_ops_by_config else ''
|
||||
exclude_unused_ops(include_ops_by_model, include_ops_by_config, get_provider_path(use_cuda=args.use_cuda))
|
||||
from exclude_unused_ops import exclude_unused_ops
|
||||
models_path = args.include_ops_by_model if args.include_ops_by_model else ''
|
||||
config_path = args.include_ops_by_config if args.include_ops_by_config else ''
|
||||
exclude_unused_ops(models_path, config_path, use_cuda=args.use_cuda)
|
||||
|
||||
if args.use_tensorrt:
|
||||
args.use_cuda = True
|
||||
|
|
|
|||
|
|
@ -15,6 +15,8 @@ from logger import log
|
|||
domain_map = {'': 'kOnnxDomain',
|
||||
'ai.onnx': 'kOnnxDomain',
|
||||
'ai.onnx.ml': 'kMLDomain',
|
||||
'ai.onnx.training': 'ai.onnx.training', # we don't have a constant for the training domains currently
|
||||
'ai.onnx.preview.training': 'ai.onnx.preview.training',
|
||||
'com.microsoft': 'kMSDomain',
|
||||
'com.microsoft.nchwc': 'kMSNchwcDomain',
|
||||
'com.microsoft.mlfeaturizers': 'kMSFeaturizersDomain',
|
||||
|
|
@ -23,67 +25,73 @@ domain_map = {'': 'kOnnxDomain',
|
|||
'com.xilinx': 'kVitisAIDomain'}
|
||||
|
||||
|
||||
def map_domain(domain):
|
||||
def _map_domain(domain):
|
||||
|
||||
if domain in domain_map:
|
||||
return domain_map[domain]
|
||||
|
||||
log.warning("Attempt to map unknown domain of {}".format(domain))
|
||||
return 'UnknownDomain'
|
||||
|
||||
|
||||
def extract_ops_from_config(file_path, referred_ops):
|
||||
def _extract_ops_from_config(file_path, required_ops):
|
||||
'''extract ops from config file of format: domain;opset;op1,op2...'''
|
||||
|
||||
if not file_path:
|
||||
return referred_ops
|
||||
return required_ops
|
||||
|
||||
if not os.path.isfile(file_path):
|
||||
log.warning('File {} does not exist'.format(file_path))
|
||||
return referred_ops
|
||||
# exit. to continue may result in unexpectedly disabling all kernels.
|
||||
log.error('Configuration file {} does not exist'.format(file_path))
|
||||
sys.exit(-1)
|
||||
|
||||
with open(file_path, 'r') as file_to_read:
|
||||
|
||||
for stripped_line in [line.strip() for line in
|
||||
file_to_read.readlines()]:
|
||||
|
||||
if not stripped_line: # skip empty lines
|
||||
continue
|
||||
|
||||
if stripped_line.startswith("#"): # skip comments
|
||||
continue
|
||||
|
||||
raw_domain, raw_opset, raw_ops =\
|
||||
[segment.strip() for segment in stripped_line.split(';')]
|
||||
|
||||
domain = map_domain(raw_domain)
|
||||
domain = _map_domain(raw_domain)
|
||||
opset = int(raw_opset)
|
||||
operators = set([raw_op.strip() for raw_op in raw_ops.split(',')])
|
||||
|
||||
if domain not in referred_ops:
|
||||
referred_ops[domain] = {opset: operators}
|
||||
if domain not in required_ops:
|
||||
required_ops[domain] = {opset: operators}
|
||||
|
||||
elif opset not in referred_ops[domain]:
|
||||
referred_ops[domain][opset] = operators
|
||||
elif opset not in required_ops[domain]:
|
||||
required_ops[domain][opset] = operators
|
||||
|
||||
else:
|
||||
referred_ops[domain][opset].update(operators)
|
||||
required_ops[domain][opset].update(operators)
|
||||
|
||||
return referred_ops # end of extract_ops_from_file(...)
|
||||
return required_ops # end of extract_ops_from_file(...)
|
||||
|
||||
|
||||
def extract_ops_from_model(model_path, referred_ops):
|
||||
def _extract_ops_from_model(model_path, required_ops):
|
||||
'''extract ops from models under model_path and return a diction'''
|
||||
|
||||
if not model_path:
|
||||
return referred_ops
|
||||
return required_ops
|
||||
|
||||
if not os.path.isdir(model_path):
|
||||
log.warning('Directory {} does not exist'.format(model_path))
|
||||
return referred_ops
|
||||
# exit. to continue may result in unexpectedly disabling all kernels.
|
||||
log.error('Directory containing models {} does not exist'.format(model_path))
|
||||
sys.exit(-1)
|
||||
|
||||
def extract_ops_from_graph(graph, operators, domain_opset_map):
|
||||
'''extract ops from graph and all subgraphs'''
|
||||
|
||||
for operator in graph.node:
|
||||
|
||||
mapped_domain = map_domain(operator.domain)
|
||||
mapped_domain = _map_domain(operator.domain)
|
||||
|
||||
if mapped_domain not in operators or\
|
||||
mapped_domain not in domain_opset_map:
|
||||
|
|
@ -110,29 +118,21 @@ def extract_ops_from_model(model_path, referred_ops):
|
|||
|
||||
for opset in model.opset_import:
|
||||
|
||||
mapped_domain = map_domain(opset.domain)
|
||||
mapped_domain = _map_domain(opset.domain)
|
||||
domain_opset_map[mapped_domain] = opset.version
|
||||
|
||||
if mapped_domain not in referred_ops:
|
||||
referred_ops[mapped_domain] = {opset.version: set()}
|
||||
if mapped_domain not in required_ops:
|
||||
required_ops[mapped_domain] = {opset.version: set()}
|
||||
|
||||
elif opset.version not in referred_ops[mapped_domain]:
|
||||
referred_ops[mapped_domain][opset.version] = set()
|
||||
elif opset.version not in required_ops[mapped_domain]:
|
||||
required_ops[mapped_domain][opset.version] = set()
|
||||
|
||||
extract_ops_from_graph(model.graph, referred_ops, domain_opset_map)
|
||||
extract_ops_from_graph(model.graph, required_ops, domain_opset_map)
|
||||
|
||||
return referred_ops # end of extract_ops_from_model(...)
|
||||
return required_ops # end of extract_ops_from_model(...)
|
||||
|
||||
|
||||
def exclude_unused_ops(model_path, config_path, provider_paths):
|
||||
'''rewrite multiple provider files'''
|
||||
|
||||
operators = extract_ops_from_config(config_path, extract_ops_from_model(model_path, {}))
|
||||
for provider_path in provider_paths:
|
||||
exclude_unused_ops_in_provider(operators, provider_path)
|
||||
|
||||
|
||||
def exclude_unused_ops_in_provider(operators, provider_path):
|
||||
def _exclude_unused_ops_in_provider(operators, provider_path):
|
||||
'''rewrite provider file to exclude unused ops'''
|
||||
|
||||
if not os.path.isfile(provider_path):
|
||||
|
|
@ -279,7 +279,14 @@ def exclude_unused_ops_in_provider(operators, provider_path):
|
|||
# end of rewrite_cpu_provider(...)
|
||||
|
||||
|
||||
def get_provider_path(ort_root=None, use_cuda=False):
|
||||
def _exclude_unused_ops_in_providers(required_operators, provider_paths):
|
||||
'''rewrite multiple provider files'''
|
||||
|
||||
for provider_path in provider_paths:
|
||||
_exclude_unused_ops_in_provider(required_operators, provider_path)
|
||||
|
||||
|
||||
def _get_provider_paths(ort_root=None, use_cuda=False):
|
||||
'''return paths to cpu and cuda providers'''
|
||||
|
||||
if not ort_root:
|
||||
|
|
@ -297,30 +304,77 @@ def get_provider_path(ort_root=None, use_cuda=False):
|
|||
return provider_paths # end of get_provider_paths
|
||||
|
||||
|
||||
def _create_config_file_with_required_ops(required_operators, model_path, config_path, output_file):
|
||||
|
||||
directory, filename = os.path.split(output_file)
|
||||
if not filename:
|
||||
log.error("Invalid path to write final config to. {}".format(output_file))
|
||||
sys.exit(-1)
|
||||
|
||||
if not os.path.exists(directory):
|
||||
os.makedirs(directory)
|
||||
|
||||
with open(output_file, 'w') as out:
|
||||
model_path_info = '--model_path {} '.format(model_path) if model_path else ''
|
||||
config_path_info = '--config_path {}'.format(config_path) if config_path else ''
|
||||
out.write("# Generated from {}{}\n".format(model_path_info, config_path_info))
|
||||
|
||||
for domain in sorted(required_operators.keys()):
|
||||
if domain == 'UnknownDomain':
|
||||
continue
|
||||
|
||||
# reverse the mapping of the domain. entry must exist given we created the required_operators dictionary.
|
||||
# also need to handle ai.onnx being special-cased as an empty string
|
||||
orig_domain = [key for (key, value) in domain_map.items() if value == domain][0]
|
||||
if not orig_domain:
|
||||
orig_domain = 'ai.onnx'
|
||||
|
||||
for opset in sorted(required_operators[domain].keys()):
|
||||
ops = required_operators[domain][opset]
|
||||
if ops:
|
||||
out.write("{};{};{}\n".format(orig_domain, opset, ','.join(sorted(ops))))
|
||||
|
||||
log.info("Wrote set of required operators to {}".format(output_file))
|
||||
|
||||
|
||||
def exclude_unused_ops(models_path, config_path, ort_root=None, use_cuda=True):
|
||||
"Note that this called directly from build.py"
|
||||
|
||||
required_operators = _extract_ops_from_config(config_path, _extract_ops_from_model(models_path, {}))
|
||||
_exclude_unused_ops_in_providers(required_operators, _get_provider_paths(ort_root, use_cuda))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Script to exclude unused operator kernels by disabling their registration in ONNXRuntime."
|
||||
"See /docs/Reduced_Operator_Kernel_build.md for more information",
|
||||
usage="Provide either model_path or config_path, or both.")
|
||||
"See /docs/Reduced_Operator_Kernel_build.md for more information",
|
||||
usage="Provide model_path, config_path, or both, to disable kernel registration of unused kernels.")
|
||||
|
||||
parser.add_argument(
|
||||
"--model_path", type=str, help="path to folder containing one or more ONNX models")
|
||||
"--model_path", type=str, help="Path to folder containing one or more ONNX models")
|
||||
|
||||
parser.add_argument(
|
||||
"--config_path", type=str, help="path to configuration file with format of 'domain;opset;op1,op2...'")
|
||||
"--config_path", type=str, help="Path to configuration file with format of 'domain;opset;op1,op2...'")
|
||||
|
||||
parser.add_argument(
|
||||
"--ort_root", type=str, help="path to ONNXRuntime repository root. "
|
||||
"--ort_root", type=str, help="Path to ONNXRuntime repository root. "
|
||||
"Inferred from the location of this script if not provided.")
|
||||
|
||||
parser.add_argument(
|
||||
"--write_combined_config_to", type=str,
|
||||
help="Optional path to create a configuration file with the combined set of required kernels "
|
||||
"from processing --model_path and/or --config_path. If provided, a configuration file will be created "
|
||||
"and NO updates will be made to the kernel registrations."
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
model_path = os.path.abspath(args.model_path) if args.model_path else ''
|
||||
models_path = os.path.abspath(args.model_path) if args.model_path else ''
|
||||
config_path = os.path.abspath(args.config_path) if args.config_path else ''
|
||||
ort_root = os.path.abspath(args.ort_root) if args.ort_root else ''
|
||||
|
||||
if not model_path and not config_path:
|
||||
if not models_path and not config_path:
|
||||
log.error('Please specify at least either model path or file path.')
|
||||
parser.print_help()
|
||||
sys.exit(-1)
|
||||
|
|
@ -328,5 +382,8 @@ if __name__ == "__main__":
|
|||
if not ort_root:
|
||||
log.info('ort_root was not specified. Inferring ORT root from location of this script.')
|
||||
|
||||
exclude_unused_ops(model_path, config_path,
|
||||
get_provider_path(ort_root, use_cuda=True))
|
||||
if not args.write_combined_config_to:
|
||||
exclude_unused_ops(models_path, config_path, ort_root, use_cuda=True)
|
||||
else:
|
||||
required_ops = _extract_ops_from_config(config_path, _extract_ops_from_model(models_path, {}))
|
||||
_create_config_file_with_required_ops(required_ops, models_path, config_path, args.write_combined_config_to)
|
||||
|
|
|
|||
Loading…
Reference in a new issue