(model, calibration_data)
| 65 | |
| 66 | |
| 67 | def quantize_model(model, calibration_data): |
| 68 | quantizer = CortexMQuantizer() |
| 69 | example_input = calibration_data[0] |
| 70 | |
| 71 | exported = torch.export.export(model, (example_input,)) |
| 72 | graph_module = exported.module() |
| 73 | |
| 74 | prepared = prepare_pt2e(graph_module, quantizer) |
| 75 | |
| 76 | logger.info(f"Calibrating with {len(calibration_data)} samples...") |
| 77 | with torch.no_grad(): |
| 78 | for i, data in enumerate(calibration_data): |
| 79 | prepared(data) |
| 80 | if (i + 1) % 25 == 0: |
| 81 | logger.info(f" Calibrated {i + 1}/{len(calibration_data)} samples") |
| 82 | |
| 83 | quantized = convert_pt2e(prepared) |
| 84 | return quantized, example_input |
| 85 | |
| 86 | |
| 87 | def export_to_pte(quantized_model, example_input, output_path: str): |
no test coverage detected