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