()
| 979 | |
| 980 | |
| 981 | def main() -> None: # noqa: C901 |
| 982 | args = _get_args() |
| 983 | |
| 984 | # Pick model from one of the supported lists |
| 985 | original_model, example_inputs = get_model_and_inputs_from_name( |
| 986 | args.model_name, args.model_input |
| 987 | ) |
| 988 | calibration_samples = load_calibration_samples( |
| 989 | args.calibration_data, example_inputs |
| 990 | ) |
| 991 | model = original_model.eval() |
| 992 | |
| 993 | # export under the assumption we quantize, the exported form also works |
| 994 | # in to_edge if we don't quantize |
| 995 | exported_program = torch.export.export( |
| 996 | model, example_inputs, strict=args.strict_export |
| 997 | ) |
| 998 | |
| 999 | model = exported_program.module() |
| 1000 | |
| 1001 | if args.enable_qdq_fusion_pass: |
| 1002 | logging.warning( |
| 1003 | "--enable_qdq_fusion_pass is deprecated and has no effect. " |
| 1004 | "Quantized node replacement is now handled within the " |
| 1005 | "respective compilation paths." |
| 1006 | ) |
| 1007 | |
| 1008 | model_name = os.path.basename(os.path.splitext(args.model_name)[0]) |
| 1009 | if args.intermediates: |
| 1010 | os.makedirs(args.intermediates, exist_ok=True) |
| 1011 | |
| 1012 | # We only support Python3.10 and above, so use a later pickle protocol |
| 1013 | torch.export.save( |
| 1014 | exported_program, |
| 1015 | f"{args.intermediates}/{model_name}_exported_program.pt2", |
| 1016 | pickle_protocol=5, |
| 1017 | ) |
| 1018 | |
| 1019 | # Quantize if required |
| 1020 | model_quant = None |
| 1021 | if args.quantize: |
| 1022 | quant_mode = QuantMode.A16W8 if "int16" in args.target else QuantMode.INT8 |
| 1023 | else: |
| 1024 | quant_mode = None |
| 1025 | |
| 1026 | if args.target.startswith("cortex-m"): |
| 1027 | # Cortex-M path: CMSIS-NN portable kernels, no delegation |
| 1028 | target_config = CortexMTargetConfig.from_target_string(args.target) |
| 1029 | if args.delegate: |
| 1030 | logging.warning( |
| 1031 | f"--delegate is ignored for target {args.target!r} " |
| 1032 | "(this target does not use delegated ops)." |
| 1033 | ) |
| 1034 | args.delegate = False |
| 1035 | model_quant, edge = _to_edge_cortex_m( |
| 1036 | exported_program, |
| 1037 | args, |
| 1038 | model, |
no test coverage detected