From fe0e676f5b3cb51d7980b5ba1d4863940dee8a82 Mon Sep 17 00:00:00 2001
From: Chen Lai <chenlai@meta.com>
Date: Sat, 28 Sep 2024 00:33:49 -0700
Subject: [PATCH] Add quantize option to the coreml script (#5710)

Summary:
Pull Request resolved: https://github.com/pytorch/executorch/pull/5710

Add the quantizer to the export script. The commandline is
```
python executorch.examples.apple.coreml.scripts.export -m resnet18 --quantize
```
bypass-github-export-checks

Reviewed By: shoumikhin

Differential Revision: D63476611

fbshipit-source-id: 480d9cd4d21e2a36d81c5a156c4869660a9259f7
---
 examples/apple/coreml/scripts/export.py | 38 +++++++++++++++++++++++--
 1 file changed, 35 insertions(+), 3 deletions(-)

diff --git a/examples/apple/coreml/scripts/export.py b/examples/apple/coreml/scripts/export.py
index 1aa5806e37..d1d4a4bcca 100644
--- a/examples/apple/coreml/scripts/export.py
+++ b/examples/apple/coreml/scripts/export.py
@@ -13,14 +13,20 @@
 import executorch.exir as exir
 
 import torch
+from coremltools.optimize.torch.quantization.quantization_config import (
+    LinearQuantizerConfig,
+    QuantizationScheme,
+)
 
 from executorch.backends.apple.coreml.compiler import CoreMLBackend
 
 from executorch.backends.apple.coreml.partition import CoreMLPartitioner
+from executorch.backends.apple.coreml.quantizer import CoreMLQuantizer
 from executorch.devtools.etrecord import generate_etrecord
 from executorch.exir import to_edge
 
 from executorch.exir.backend.backend_api import to_backend
+from torch.ao.quantization.quantize_pt2e import convert_pt2e, prepare_pt2e
 
 from torch.export import export
 
@@ -75,6 +81,13 @@ def parse_args() -> argparse.ArgumentParser:
     parser.add_argument("--generate_etrecord", action=argparse.BooleanOptionalAction)
     parser.add_argument("--save_processed_bytes", action=argparse.BooleanOptionalAction)
 
+    parser.add_argument(
+        "--quantize",
+        action=argparse.BooleanOptionalAction,
+        required=False,
+        help="Quantize CoreML model",
+    )
+
     args = parser.parse_args()
     return args
 
@@ -110,9 +123,10 @@ def export_lowered_module_to_executorch_program(lowered_module, example_inputs):
     return exec_prog
 
 
-def save_executorch_program(exec_prog, model_name, compute_unit):
+def save_executorch_program(exec_prog, model_name, compute_unit, quantize):
     buffer = exec_prog.buffer
-    filename = f"{model_name}_coreml_{compute_unit}.pte"
+    data_type = "quantize" if quantize else "fp"
+    filename = f"{model_name}_coreml_{compute_unit}_{data_type}.pte"
     print(f"Saving exported program to {filename}")
     with open(filename, "wb") as file:
         file.write(buffer)
@@ -168,6 +182,22 @@ def main():
     if args.use_partitioner:
         model.eval()
         exir_program_aten = torch.export.export(model, example_inputs)
+        if args.quantize:
+            quantization_config = LinearQuantizerConfig.from_dict(
+                {
+                    "global_config": {
+                        "quantization_scheme": QuantizationScheme.affine,
+                        "activation_dtype": torch.quint8,
+                        "weight_dtype": torch.qint8,
+                        "weight_per_channel": True,
+                    }
+                }
+            )
+
+            quantizer = CoreMLQuantizer(quantization_config)
+            model = prepare_pt2e(model, quantizer)  # pyre-fixme[6]
+            model(*example_inputs)
+            exir_program_aten = convert_pt2e(model)
 
         edge_program_manager = exir.to_edge(exir_program_aten)
         edge_copy = copy.deepcopy(edge_program_manager)
@@ -189,7 +219,9 @@ def main():
             example_inputs,
         )
 
-    save_executorch_program(exec_program, args.model_name, args.compute_unit)
+    save_executorch_program(
+        exec_program, args.model_name, args.compute_unit, args.quantize
+    )
     generate_etrecord(f"{args.model_name}_coreml_etrecord.bin", edge_copy, exec_program)
 
     if args.save_processed_bytes and lowered_module is not None: