Function to_quantized_edge_program

backends/nxp/tests/executorch_pipeline.py:178–267

(
    model: torch.nn.Module,
    input_spec: Iterable[ModelInputSpec] | tuple[int, ...] | list[tuple[int, ...]],
    operators_not_to_delegate: list[str] = None,
    get_calibration_inputs_fn: GetCalibrationInputsFn = get_random_calibration_inputs,
    target: str = "imxrt700",
    use_qat: bool = False,
    train_fn: Callable[[torch.fx.GraphModule], None] | None = None,
    remove_quant_io_ops: bool = False,
    custom_delegation_options: CustomDelegationOptions = CustomDelegationOptions(),  # noqa B008
    get_quantizer_fn: Callable[[], Quantizer] | None = None,
    use_neutron_for_format_conversion: bool = True,
    use_quant_state_dict: bool = True,
    fetch_constants_to_sram: bool = False,
    dump_kernel_selection_code: bool = False,
    use_new_flow_neutron_c: bool = False,
    delegate_to_npu=True,
)

def to_quantized_edge_program(
    model: torch.nn.Module,
    input_spec: Iterable[ModelInputSpec] | tuple[int, ...] | list[tuple[int, ...]],
    operators_not_to_delegate: list[str] = None,
    get_calibration_inputs_fn: GetCalibrationInputsFn = get_random_calibration_inputs,
    target: str = "imxrt700",
    use_qat: bool = False,
    train_fn: Callable[[torch.fx.GraphModule], None] | None = None,
    remove_quant_io_ops: bool = False,
    custom_delegation_options: CustomDelegationOptions = CustomDelegationOptions(),  # noqa B008
    get_quantizer_fn: Callable[[], Quantizer] | None = None,
    use_neutron_for_format_conversion: bool = True,
    use_quant_state_dict: bool = True,
    fetch_constants_to_sram: bool = False,
    dump_kernel_selection_code: bool = False,
    use_new_flow_neutron_c: bool = False,
    delegate_to_npu=True,
) -> EdgeProgramManager:
    _neutron_target_spec = NeutronTargetSpec(target)
    custom_delegation_options.use_new_flow_neutron_c = use_new_flow_neutron_c
    if get_quantizer_fn is None:
        get_quantizer_fn = partial(
            _get_default_quantizer, _neutron_target_spec, use_qat
        )
    input_spec = to_model_input_spec(input_spec)
    calibration_inputs = get_calibration_inputs_fn(input_spec)
    example_input = _get_example_input(input_spec)

    # Make sure the model is in the evaluation mode.
    model.eval()

    exir_program_aten = torch.export.export(model, example_input, strict=True)

    exir_program_aten__module_quant = calibrate_and_quantize(
        model=exir_program_aten,
        calibration_inputs=calibration_inputs,
        quantizer=get_quantizer_fn(),
        is_qat=use_qat,
        train_fn=train_fn,
    )

    # List of operators to not decompose during the lowering.
    preserve_ops = [torch.ops.aten.prelu.default]
    compile_spec = generate_neutron_compile_spec(
        target,
        operators_not_to_delegate=operators_not_to_delegate,
        use_neutron_for_format_conversion=use_neutron_for_format_conversion,
        fetch_constants_to_sram=fetch_constants_to_sram,
        dump_kernel_selection_code=dump_kernel_selection_code,
        use_new_flow_neutron_c=use_new_flow_neutron_c,
    )
    post_quant_state_dict = (
        exir_program_aten__module_quant.state_dict() if use_quant_state_dict else None
    )
    if delegate_to_npu:
        partitioners = [
            NeutronPartitioner(
                compile_spec,