(
model: nn.Module,
example_inputs: Tuple[Any, ...],
file_name: str = "CadenceDemoModel",
working_dir: Optional[str] = None,
)
| 38 | |
| 39 | |
| 40 | def export_model( |
| 41 | model: nn.Module, |
| 42 | example_inputs: Tuple[Any, ...], |
| 43 | file_name: str = "CadenceDemoModel", |
| 44 | working_dir: Optional[str] = None, |
| 45 | ) -> ExecutorchProgramManager: |
| 46 | # create work directory for outputs and model binary |
| 47 | if working_dir is None: |
| 48 | working_dir = tempfile.mkdtemp(dir="/tmp") |
| 49 | logging.debug(f"Created work directory {working_dir}") |
| 50 | |
| 51 | # Instantiate the quantizer |
| 52 | quantizer = CadenceDefaultQuantizer() |
| 53 | |
| 54 | # Prepare the model |
| 55 | prepared_gm = prepare_pt2(model, example_inputs, quantizer) |
| 56 | |
| 57 | # Calibrate the model |
| 58 | for samples in [example_inputs]: |
| 59 | prepared_gm(*samples) |
| 60 | |
| 61 | # Convert the model |
| 62 | converted_model = convert_pt2(prepared_gm) |
| 63 | |
| 64 | # Get reference outputs from converted model |
| 65 | ref_outputs = converted_model(*example_inputs) |
| 66 | |
| 67 | ep = torch.export.export(converted_model, example_inputs, strict=True) |
| 68 | |
| 69 | # Fuse the quantized patterns on the exported program (note: quantizer needs to be the same as the one used in prepare_and_convert_pt2) |
| 70 | ep = apply_pre_edge_transform_passes(ep, quantizer) |
| 71 | |
| 72 | # Get edge program after Cadence specific passes |
| 73 | exec_prog: ExecutorchProgramManager = _lower_ep_to_cadence_gen_etrecord( |
| 74 | ep, output_dir=working_dir |
| 75 | ) |
| 76 | |
| 77 | logging.info("Final exported graph:\n") |
| 78 | exec_prog.exported_program().graph_module.graph.print_tabular() |
| 79 | |
| 80 | forward_test_data = BundledProgramManager.bundled_program_test_data_gen( |
| 81 | method="forward", inputs=example_inputs, expected_outputs=ref_outputs |
| 82 | ) |
| 83 | bundled_program_manager = BundledProgramManager([forward_test_data]) |
| 84 | buffer = bundled_program_manager._serialize( |
| 85 | exec_prog, |
| 86 | bundled_program_manager.get_method_test_suites(), |
| 87 | forward_test_data, |
| 88 | ) |
| 89 | # Save the program as pte (default name is CadenceDemoModel.pte) |
| 90 | save_pte_program(exec_prog, file_name, working_dir) |
| 91 | # Save the program as btpe (default name is CadenceDemoModel.bpte) |
| 92 | save_bpte_program(buffer, file_name, working_dir) |
| 93 | |
| 94 | logging.debug( |
| 95 | f"Executorch bundled program buffer saved to {file_name} is {len(buffer)} total bytes" |
| 96 | ) |
| 97 |
no test coverage detected