Add ability to generate configuration file with required operators. (#5089)

* Add ability to generate configuration file with required operators.
This commit is contained in:
Scott McKay 2020-09-09 21:39:17 +10:00 committed by GitHub
parent 80ada0291f
commit dbf4e7019d
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 110 additions and 47 deletions

View file

@ -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.

View file

@ -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

View file

@ -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)