(
module: torch.nn.Module,
inputs,
calibration_dataset,
precision: Precision,
is_per_channel: bool = True,
is_qat: bool = False,
)
| 53 | |
| 54 | |
| 55 | def quantize_module( |
| 56 | module: torch.nn.Module, |
| 57 | inputs, |
| 58 | calibration_dataset, |
| 59 | precision: Precision, |
| 60 | is_per_channel: bool = True, |
| 61 | is_qat: bool = False, |
| 62 | ) -> torch.nn.Module: |
| 63 | quantizer = EnnQuantizer() |
| 64 | quantizer.setup_quant_params(precision, is_per_channel, is_qat) |
| 65 | logging.info("Export nn module for quantization...") |
| 66 | exported_module = torch.export.export(module, inputs).module() |
| 67 | DecomposeScaledDotProductAttention()(exported_module) |
| 68 | logging.info("Quantizing the module...") |
| 69 | annotated_module = prepare_pt2e(exported_module, quantizer) |
| 70 | for data in calibration_dataset: |
| 71 | annotated_module(*data) |
| 72 | quantized_module = convert_pt2e(annotated_module, fold_quantize=False) |
| 73 | logging.info("Quantizing finished.") |
| 74 | return quantized_module |
| 75 | |
| 76 | |
| 77 | def to_edge_transform_and_lower_to_enn( |
no test coverage detected