ai-edge-quantizer-nightly 0.0.1.dev20241210__py3-none-any.whl → 0.0.1.dev20241218__py3-none-any.whl

ai_edge_quantizer/calibrator.py

@@ -41,6 +41,7 @@ class Calibrator:
   def __init__(
       self,
       float_tflite: Union[str, bytes],
+      num_threads: int = 16,
   ):
     self._flatbuffer_model = tfl_flatbuffer_utils.read_model(float_tflite)
 
@@ -50,7 +51,7 @@ class Calibrator:
           " the model (e.g., if it is already quantized)."
       )
     self._tfl_interpreter = tfl_interpreter_utils.create_tfl_interpreter(
-        float_tflite
+        float_tflite, use_xnnpack=True, num_threads=num_threads
     )
     # Tensor name to tensor content.
     self._tensor_content_map: dict[str, Any] = {}

ai_edge_quantizer/model_validator.py

@@ -207,7 +207,8 @@ def _setup_validation_interpreter(
     model: bytes,
     signature_input: dict[str, Any],
     signature_key: Optional[str],
-    use_reference_kernel: bool,
+    use_xnnpack: bool,
+    num_threads: int,
 ) -> tuple[Any, int, dict[str, Any]]:
   """Setup the interpreter for validation given a signature key.
 
@@ -216,15 +217,15 @@ def _setup_validation_interpreter(
     signature_input: A dictionary of input tensor name and its value.
     signature_key: The signature key to be used for invoking the models. If the
       model only has one signature, this can be set to None.
-    use_reference_kernel: Whether to use the reference kernel for the
-      interpreter.
+    use_xnnpack: Whether to use xnnpack for the interpreter.
+    num_threads: The number of threads to use for the interpreter.
 
   Returns:
     A tuple of interpreter, subgraph_index and tensor_name_to_details.
   """
 
   interpreter = utils.create_tfl_interpreter(
-      tflite_model=model, use_reference_kernel=use_reference_kernel
+      tflite_model=model, use_xnnpack=use_xnnpack, num_threads=num_threads
   )
   utils.invoke_interpreter_signature(
       interpreter, signature_input, signature_key
@@ -247,7 +248,8 @@ def compare_model(
     test_data: dict[str, Iterable[dict[str, Any]]],
     error_metric: str,
     compare_fn: Callable[[Any, Any], float],
-    use_reference_kernel: bool = False,
+    use_xnnpack: bool = True,
+    num_threads: int = 16,
 ) -> ComparisonResult:
   """Compares model tensors over a model signature using the compare_fn.
 
@@ -266,8 +268,8 @@ def compare_model(
     compare_fn: a comparison function to be used for calculating the statistics,
       this function must be taking in two ArrayLike strcuture and output a
       single float value.
-    use_reference_kernel: Whether to use the reference kernel for the
-      interpreter.
+    use_xnnpack: Whether to use xnnpack for the interpreter.
+    num_threads: The number of threads to use for the interpreter.
 
   Returns:
     A ComparisonResult object.
@@ -282,12 +284,17 @@ def compare_model(
               reference_model,
               signature_input,
               signature_key,
-              use_reference_kernel,
+              use_xnnpack,
+              num_threads,
           )
       )
       targ_interpreter, targ_subgraph_index, targ_tensor_name_to_details = (
           _setup_validation_interpreter(
-              target_model, signature_input, signature_key, use_reference_kernel
+              target_model,
+              signature_input,
+              signature_key,
+              use_xnnpack,
+              num_threads,
           )
       )
       # Compare the cached tensor values.

ai_edge_quantizer/quantizer.py

@@ -216,6 +216,7 @@ class Quantizer:
       self,
       calibration_data: dict[str, Iterable[_SignatureInput]],
       previous_calibration_result: Optional[_CalibrationResult] = None,
+      num_threads: int = 16,
   ) -> _CalibrationResult:
     """Calibrates the float model (required by static range quantization).
 
@@ -223,6 +224,7 @@ class Quantizer:
       calibration_data: Calibration data for a model signature.
       previous_calibration_result: Previous calibration result to be loaded. The
         calibration process will be resumed from the previous result.
+      num_threads: Number of threads to use for calibration.
 
     Returns:
       Calibration result ({tensor_name: tensor QSVs (e.g.,min/max)}).
@@ -233,7 +235,7 @@ class Quantizer:
     if not self.need_calibration:
       return {}
 
-    calib = calibrator.Calibrator(self.float_model)
+    calib = calibrator.Calibrator(self.float_model, num_threads=num_threads)
     if previous_calibration_result is not None:
       calib.load_model_qsvs(previous_calibration_result)
     calib.calibrate(calibration_data, self._recipe_manager)
@@ -297,7 +299,8 @@ class Quantizer:
       self,
       test_data: Optional[dict[str, Iterable[_SignatureInput]]] = None,
       error_metrics: str = 'mse',
-      use_reference_kernel: bool = False,
+      use_xnnpack: bool = True,
+      num_threads: int = 16,
   ) -> model_validator.ComparisonResult:
     """Numerical validation of the quantized model for a model signature.
 
@@ -314,7 +317,8 @@ class Quantizer:
         data that will be used for validation. If set to None, random normal
         distributed data will be used for all signatures in the model.
       error_metrics: Error metrics to be used for comparison.
-      use_reference_kernel: Whether to use the reference kernel for validation.
+      use_xnnpack: Whether to use the xnnpack library for validation.
+      num_threads: Number of threads to use for validation.
 
     Returns:
       The comparison result.
@@ -330,7 +334,8 @@ class Quantizer:
         test_data,
         error_metrics,
         validation_utils.get_validation_func(error_metrics),
-        use_reference_kernel=use_reference_kernel,
+        use_xnnpack=use_xnnpack,
+        num_threads=num_threads,
     )
 
   def _get_quantization_params(

ai_edge_quantizer/utils/tfl_interpreter_utils.py

@@ -30,15 +30,16 @@ DEFAULT_SIGNATURE_KEY = "serving_default"
 def create_tfl_interpreter(
     tflite_model: Union[str, bytes],
     allocate_tensors: bool = True,
-    use_reference_kernel: bool = False,
+    use_xnnpack: bool = True,
+    num_threads: int = 16,
 ) -> tfl.Interpreter:
   """Creates a TFLite interpreter from a model file.
 
   Args:
     tflite_model: Model file path or bytes.
     allocate_tensors: Whether to allocate tensors.
-    use_reference_kernel: Whether to use the reference kernel for the
-      interpreter.
+    use_xnnpack: Whether to use the XNNPACK delegate for the interpreter.
+    num_threads: The number of threads to use for the interpreter.
 
   Returns:
     A TFLite interpreter.
@@ -47,12 +48,13 @@ def create_tfl_interpreter(
     with gfile.GFile(tflite_model, "rb") as f:
       tflite_model = f.read()
 
-  if use_reference_kernel:
-    op_resolver = tfl.OpResolverType.BUILTIN_REF
+  if use_xnnpack:
+    op_resolver = tfl.OpResolverType.BUILTIN
   else:
     op_resolver = tfl.OpResolverType.BUILTIN_WITHOUT_DEFAULT_DELEGATES
   tflite_interpreter = tfl.Interpreter(
       model_content=bytes(tflite_model),
+      num_threads=num_threads,
       experimental_op_resolver_type=op_resolver,
       experimental_preserve_all_tensors=True,
   )

{ai_edge_quantizer_nightly-0.0.1.dev20241210.dist-info → ai_edge_quantizer_nightly-0.0.1.dev20241218.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-quantizer-nightly
-Version: 0.0.1.dev20241210
+Version: 0.0.1.dev20241218
 Summary: A quantizer for advanced developers to quantize converted AI Edge models.
 Home-page: https://github.com/google-ai-edge/ai-edge-quantizer
 Keywords: On-Device ML,AI,Google,TFLite,Quantization,LLMs,GenAI

{ai_edge_quantizer_nightly-0.0.1.dev20241210.dist-info → ai_edge_quantizer_nightly-0.0.1.dev20241218.dist-info}/RECORD

@@ -2,18 +2,18 @@ ai_edge_quantizer/__init__.py,sha256=4pFSkukSwahYyzwqia0yPRyz8TnFQfGRthVJhYpMWas
 ai_edge_quantizer/algorithm_manager.py,sha256=9nd4Txfl2z-14rFHmL7vqSfnkAQeagCRKyCIQ7ru0_Y,5981
 ai_edge_quantizer/algorithm_manager_api.py,sha256=u903TG0s1uIDhJqfeJne3CFl8A93phZrwgV2-hwdcXU,9247
 ai_edge_quantizer/algorithm_manager_api_test.py,sha256=tL_ozYFTsOPX8qGcti0KTz37nVsCxf0SSG5C45SyT-g,7319
-ai_edge_quantizer/calibrator.py,sha256=BSu0DPzVhAgFFA0JsHZtawPFlr0YPirRxItuA9SerNg,11007
+ai_edge_quantizer/calibrator.py,sha256=0zAWrSpl_08u6BNLVgG_TQeNcT16wJ-oLeQgznziGoo,11079
 ai_edge_quantizer/calibrator_test.py,sha256=5DGvKWRRjjU3L5wZoN56AyOVljmxOitwhuBUp6GL_bU,11354
 ai_edge_quantizer/conftest.py,sha256=SxCz-5LlRD_lQm4hQc4c6IGG7DS8d7IyEWY9gnscPN0,794
 ai_edge_quantizer/default_policy.py,sha256=TQ9yY8jtrSpMsTBsTyKW6TY-voGH_psvwGZoFglAbiA,9079
 ai_edge_quantizer/model_modifier.py,sha256=Z8EYtrz4zhCFpzd1zVwl2AetVE3BGBf5OvB2DbVQuds,5850
 ai_edge_quantizer/model_modifier_test.py,sha256=cJd04SLOG-fQZZNZPcisoBLx3cLtWEwGqUBbLb-pif4,4751
-ai_edge_quantizer/model_validator.py,sha256=QvlG1TewSBo9FMwzDYPFGqR4mOa_Xhn21wi2OFAvbCI,12593
+ai_edge_quantizer/model_validator.py,sha256=oZk0b1qGczaEm5erJFm4SbwadDnl7DFhC0bXuxwVgps,12787
 ai_edge_quantizer/model_validator_test.py,sha256=ctvVmMHvnmFbkG4o8Jaa6kXXRrGHzhYpNylgLSmOboA,12951
 ai_edge_quantizer/params_generator.py,sha256=FvBub5yM2q98k7wNLgEyRerf8sVIETvGbrFcXFPUPdA,13523
 ai_edge_quantizer/params_generator_test.py,sha256=d9JwR-yxNJgg1SW-m8sFFPkIRdhgsDwMpVKsBQFL0gg,37658
 ai_edge_quantizer/qtyping.py,sha256=bue_WfK05QTkQcoyVVWeIxh8LRVGhHMWruXk3cgpFpw,14577
-ai_edge_quantizer/quantizer.py,sha256=OYfSo06JcoursXbJBRfHQbR2-Pa4sHnZB2n9od9OzEY,13039
+ai_edge_quantizer/quantizer.py,sha256=Gny7WLuRibiIuDtcRn_g8RCD-zAm_fuDG7WmGq5dRx8,13238
 ai_edge_quantizer/quantizer_test.py,sha256=38oTMJwMmxwPDeqT3eaVbazjtuIUIzMQ3mJNKh_eNQY,20493
 ai_edge_quantizer/recipe.py,sha256=r5tJiUs-ihZFzeK_jP2sUIUgTqZsL5SWvbUokuIUPDo,2251
 ai_edge_quantizer/recipe_manager.py,sha256=qcGUD7e7BISKdsY9WH2rdaRR3acmzSA5qMezGNbzlpo,8931
@@ -52,12 +52,12 @@ ai_edge_quantizer/utils/calibration_utils_test.py,sha256=Z-AcdTieesWFKyKBb08ZXm4
 ai_edge_quantizer/utils/test_utils.py,sha256=95BDAdjE4Zvd6JZ90fG8FE3wKWE-Lu0ZIE3hQ1B6adI,3616
 ai_edge_quantizer/utils/tfl_flatbuffer_utils.py,sha256=F6_AkCSv35FAhJX2qel8VTARhGOVwaeo7_mqRZygrpA,10126
 ai_edge_quantizer/utils/tfl_flatbuffer_utils_test.py,sha256=AbyDxoM62k4ojD8gPdkWo--xe5hlX3t0kobQSA80kuk,7740
-ai_edge_quantizer/utils/tfl_interpreter_utils.py,sha256=GzrsaL3fkOXN5iPRJv7lqhNISY6lnrBVTotWDHzI5m8,10344
+ai_edge_quantizer/utils/tfl_interpreter_utils.py,sha256=yta7b_VmhVZmntwHK27vqVnie3XRejN459P0uJHbpb8,10431
 ai_edge_quantizer/utils/tfl_interpreter_utils_test.py,sha256=Op3JxtOqlrjzmYF18jnnstL1k9xiY9kKJ8S2vklKGkc,11327
 ai_edge_quantizer/utils/validation_utils.py,sha256=oYw33Sg547AqtGw-choPUJmp9SAKkV46J_ddqSsum2Q,3950
 ai_edge_quantizer/utils/validation_utils_test.py,sha256=V_qNDikPD4OPB-siOLQCWNVWTAu87h2IgNYt7teFd-o,2934
-ai_edge_quantizer_nightly-0.0.1.dev20241210.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-ai_edge_quantizer_nightly-0.0.1.dev20241210.dist-info/METADATA,sha256=o1xW7CHGdW5K9XN9eUik8kRTwfVhjiKMJijm4Sewl4M,1484
-ai_edge_quantizer_nightly-0.0.1.dev20241210.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_quantizer_nightly-0.0.1.dev20241210.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
-ai_edge_quantizer_nightly-0.0.1.dev20241210.dist-info/RECORD,,
+ai_edge_quantizer_nightly-0.0.1.dev20241218.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+ai_edge_quantizer_nightly-0.0.1.dev20241218.dist-info/METADATA,sha256=9YqAIud-Y9td1FerL7kzgo_dyJh7qiBR0f6tEy_aHtg,1484
+ai_edge_quantizer_nightly-0.0.1.dev20241218.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_quantizer_nightly-0.0.1.dev20241218.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
+ai_edge_quantizer_nightly-0.0.1.dev20241218.dist-info/RECORD,,