def convert_to_onnx(reporter, model, output_dir, args, template_variables):
reporter.print_section_heading('{}Converting {} to ONNX',
'(DRY RUN) ' if args.dry_run else '',
model.name)
converter_path = Path(__file__).absolute().parent / \
'internal_scripts' / model.converter_to_onnx
conversion_to_onnx_args = [
string.Template(arg).substitute(template_variables)
for arg in model.conversion_to_onnx_args
]
cmd = [
str(args.python), '--',
str(converter_path), *conversion_to_onnx_args
]
reporter.print('Conversion to ONNX command: {}',
_common.command_string(cmd))
reporter.print(flush=True)
success = True if args.dry_run else reporter.job_context.subprocess(cmd)
reporter.print()
return success
def run_pre_convert(reporter, model, output_dir, args):
script = _common.MODEL_ROOT / model.subdirectory / 'pre-convert.py'
if not script.exists():
return True
reporter.print_section_heading('{}Running pre-convert script for {}',
'(DRY RUN) ' if args.dry_run else '', model.name)
cmd = [str(args.python), '--', str(script), '--',
str(args.download_dir / model.subdirectory), str(output_dir / model.subdirectory)]
reporter.print('Pre-convert command: {}', _common.command_string(cmd))
reporter.print(flush=True)
success = True if args.dry_run else reporter.job_context.subprocess(cmd)
reporter.print()
return success
def quantize(self, reporter, model, precision, target_device,
pot_cmd_prefix, pot_env) -> bool:
input_precision = _common.KNOWN_QUANTIZED_PRECISIONS[precision]
pot_config_base_path = _common.MODEL_ROOT / model.subdirectory / 'quantization.yml'
try:
with pot_config_base_path.open('rb') as pot_config_base_file:
pot_config_base = yaml.safe_load(pot_config_base_file)
except FileNotFoundError:
pot_config_base = DEFAULT_POT_CONFIG_BASE
pot_config_paths = {
'engine': {
'config':
str(_common.MODEL_ROOT / model.subdirectory /
'accuracy-check.yml'),
},
'model': {
'model':
str(self.model_dir / model.subdirectory / input_precision /
(model.name + '.xml')),
'weights':
str(self.model_dir / model.subdirectory / input_precision /
(model.name + '.bin')),
'model_name':
model.name,
}
}
pot_config = {**pot_config_base, **pot_config_paths}
if target_device:
pot_config['compression']['target_device'] = target_device
reporter.print_section_heading('{}Quantizing {} from {} to {}',
'(DRY RUN) ' if self.dry_run else '',
model.name, input_precision, precision)
model_output_dir = self.output_dir / model.subdirectory / precision
pot_config_path = model_output_dir / 'pot-config.json'
reporter.print('Creating {}...', pot_config_path)
pot_config_path.parent.mkdir(parents=True, exist_ok=True)
with pot_config_path.open('w') as pot_config_file:
json.dump(pot_config, pot_config_file, indent=4)
pot_config_file.write('\n')
pot_output_dir = model_output_dir / 'pot-output'
pot_output_dir.mkdir(parents=True, exist_ok=True)
pot_cmd = [
*pot_cmd_prefix,
'--config={}'.format(pot_config_path),
'--direct-dump',
'--output-dir={}'.format(pot_output_dir),
]
reporter.print('Quantization command: {}',
_common.command_string(pot_cmd))
reporter.print(
'Quantization environment: {}',
' '.join('{}={}'.format(k, _common.quote_arg(v))
for k, v in sorted(pot_env.items())))
success = True
if not self.dry_run:
reporter.print(flush=True)
success = reporter.job_context.subprocess(pot_cmd,
env={
**os.environ,
**pot_env
})
reporter.print()
if not success: return False
if not self.dry_run:
reporter.print('Moving quantized model to {}...', model_output_dir)
for ext in ['.xml', '.bin']:
(pot_output_dir / 'optimized' / (model.name + ext)).replace(
model_output_dir / (model.name + ext))
reporter.print()
return True
def convert(reporter, model, output_dir, args, mo_props, requested_precisions):
telemetry = _common.Telemetry()
if model.mo_args is None:
reporter.print_section_heading('Skipping {} (no conversions defined)',
model.name)
reporter.print()
return True
model_precisions = requested_precisions & model.precisions
if not model_precisions:
reporter.print_section_heading('Skipping {} (all conversions skipped)',
model.name)
reporter.print()
return True
(output_dir / model.subdirectory).mkdir(parents=True, exist_ok=True)
if not run_pre_convert(reporter, model, output_dir, args):
telemetry.send_event('md', 'converter_failed_models', model.name)
telemetry.send_event(
'md', 'converter_error',
json.dumps({
'error': 'pre-convert-script-failed',
'model': model.name,
'precision': None
}))
return False
model_format = model.framework
mo_extension_dir = mo_props.base_dir / 'extensions'
if not mo_extension_dir.exists():
mo_extension_dir = mo_props.base_dir
template_variables = {
'config_dir': _common.MODEL_ROOT / model.subdirectory,
'conv_dir': output_dir / model.subdirectory,
'dl_dir': args.download_dir / model.subdirectory,
'mo_dir': mo_props.base_dir,
'mo_ext_dir': mo_extension_dir,
}
if model.conversion_to_onnx_args:
if not convert_to_onnx(reporter, model, output_dir, args,
template_variables):
telemetry.send_event('md', 'converter_failed_models', model.name)
telemetry.send_event(
'md', 'converter_error',
json.dumps({
'error': 'convert_to_onnx-failed',
'model': model.name,
'precision': None
}))
return False
model_format = 'onnx'
expanded_mo_args = [
string.Template(arg).substitute(template_variables)
for arg in model.mo_args
]
for model_precision in sorted(model_precisions):
data_type = model_precision.split('-')[0]
layout_string = ','.join('{}({})'.format(input.name, input.layout)
for input in model.input_info if input.layout)
shape_string = ','.join(
str(input.shape) for input in model.input_info if input.shape)
if layout_string:
expanded_mo_args.append('--layout={}'.format(layout_string))
if shape_string:
expanded_mo_args.append('--input_shape={}'.format(shape_string))
mo_cmd = [
*mo_props.cmd_prefix, '--framework={}'.format(model_format),
'--data_type={}'.format(data_type), '--output_dir={}'.format(
output_dir / model.subdirectory / model_precision),
'--model_name={}'.format(model.name),
'--input={}'.format(','.join(input.name
for input in model.input_info)),
*expanded_mo_args, *mo_props.extra_args
]
reporter.print_section_heading('{}Converting {} to IR ({})',
'(DRY RUN) ' if args.dry_run else '',
model.name, model_precision)
reporter.print('Conversion command: {}',
_common.command_string(mo_cmd))
if not args.dry_run:
reporter.print(flush=True)
if not reporter.job_context.subprocess(mo_cmd):
telemetry.send_event('md', 'converter_failed_models',
model.name)
telemetry.send_event(
'md', 'converter_error',
json.dumps({
'error': 'mo-failed',
'model': model.name,
'precision': model_precision
}))
return False
reporter.print()
return True
