ai-edge-torch-nightly 0.2.0.dev20240714__py3-none-any.whl → 0.3.0.dev20240926__py3-none-any.whl

ai_edge_torch/__init__.py

@@ -13,10 +13,11 @@
 # limitations under the License.
 # ==============================================================================
 
-from .convert.converter import convert
-from .convert.converter import signature
-from .convert.to_channel_last_io import to_channel_last_io
-from .model import Model
+from ai_edge_torch._convert.converter import convert
+from ai_edge_torch._convert.converter import signature
+from ai_edge_torch._convert.to_channel_last_io import to_channel_last_io
+from ai_edge_torch.model import Model
+from ai_edge_torch.version import __version__
 
 
 def load(path: str) -> Model:

ai_edge_torch/_convert/conversion.py

@@ -0,0 +1,112 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import logging
+import os
+from typing import Any, Optional
+
+from ai_edge_torch import fx_pass_base
+from ai_edge_torch import lowertools
+from ai_edge_torch import model
+from ai_edge_torch._convert import fx_passes
+from ai_edge_torch._convert import signature
+from ai_edge_torch.generative import fx_passes as generative_fx_passes
+from ai_edge_torch.quantize import quant_config as qcfg
+import torch
+
+os.environ["EXPERIMENTAL_XLA_UNBOUNDED_DYNAMISM"] = "1"
+
+
+def _run_convert_passes(
+    exported_program: torch.export.ExportedProgram,
+) -> torch.export.ExportedProgram:
+  exported_program = generative_fx_passes.run_generative_passes(
+      exported_program
+  )
+  return fx_pass_base.run_passes(
+      exported_program,
+      [
+          fx_passes.BuildInterpolateCompositePass(),
+          fx_passes.CanonicalizePass(),
+          fx_passes.OptimizeLayoutTransposesPass(),
+          fx_passes.CanonicalizePass(),
+          fx_passes.BuildAtenCompositePass(),
+          fx_passes.CanonicalizePass(),
+          fx_passes.InjectMlirDebuginfoPass(),
+          fx_passes.CanonicalizePass(),
+      ],
+  )
+
+
+def _warn_training_modules(signatures: list[signature.Signature]):
+  """Warns the user if the module is in training mode (.eval not called)."""
+  for sig in signatures:
+    if not sig.module.training:
+      continue
+
+    message = (
+        "Your model {sig_name}is converted in training mode. Please set the"
+        " module in evaluation mode with `module.eval()` for better on-device"
+        " performance and compatibility."
+    )
+    if len(signatures) == 1 and sig.name == model.DEFAULT_SIGNATURE_NAME:
+      # User does not specify any signature names explicitly.
+      message = message.format(sig_name="")
+    else:
+      message = message.format(sig_name=f'"{sig.name}" ')
+
+    logging.warning(message)
+
+
+def convert_signatures(
+    signatures: list[signature.Signature],
+    *,
+    quant_config: Optional[qcfg.QuantConfig] = None,
+    _tfl_converter_flags: Optional[dict[str, Any]],
+) -> model.TfLiteModel:
+  """Converts a list of `signature.Signature`s and embeds them into one `model.TfLiteModel`.
+
+  Args:
+      signatures: The list of 'signature.Signature' objects containing PyTorch
+        modules to be converted.
+      quant_config: User-defined quantization method and scheme of the model.
+      _tfl_converter_flags: A nested dictionary allowing setting flags for the
+        underlying tflite converter.
+
+  Returns:
+    The converted `model.TfLiteModel` object.
+  """
+  if _tfl_converter_flags is None:
+    _tfl_converter_flags = {}
+
+  _warn_training_modules(signatures)
+
+  exported_programs: torch.export.torch.export.ExportedProgram = [
+      torch.export.export(
+          sig.module, sig.flat_args, dynamic_shapes=sig.dynamic_shapes
+      )
+      for sig in signatures
+  ]
+
+  # Apply default fx passes
+  exported_programs = list(map(_run_convert_passes, exported_programs))
+  tflite_model = lowertools.exported_programs_to_tflite(
+      exported_programs,
+      signatures,
+      quant_config=quant_config,
+      _tfl_converter_flags=_tfl_converter_flags,
+  )
+
+  return model.TfLiteModel(tflite_model)

ai_edge_torch/_convert/conversion_utils.py

@@ -0,0 +1,64 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from typing import Any
+
+from ai_edge_torch.quantize import quant_config as qcfg
+import tensorflow as tf
+
+
+def apply_tfl_converter_flags(
+    converter: tf.lite.TFLiteConverter, tfl_converter_flags: dict[str, Any]
+):
+  """Applies TFLite converter flags to the converter.
+
+  Args:
+    converter: TFLite converter.
+    tfl_converter_flags: TFLite converter flags.
+  """
+
+  def _set_converter_flag(path: list[Any]):
+    if len(path) < 2:
+      raise ValueError("Expecting at least two values in the path.")
+
+    target_obj = converter
+    for idx in range(len(path) - 2):
+      target_obj = getattr(target_obj, path[idx])
+
+    setattr(target_obj, path[-2], path[-1])
+
+  def _iterate_dict_tree(flags_dict: dict[str, Any], path: list[Any]):
+    for key, value in flags_dict.items():
+      path.append(key)
+      if isinstance(value, dict):
+        _iterate_dict_tree(value, path)
+      else:
+        path.append(value)
+        _set_converter_flag(path)
+        path.pop()
+      path.pop()
+
+  _iterate_dict_tree(tfl_converter_flags, [])
+
+
+def set_tfl_converter_quant_flags(
+    converter: tf.lite.TFLiteConverter, quant_config: qcfg.QuantConfig
+):
+  if quant_config is not None:
+    quantizer_mode = quant_config._quantizer_mode
+    if quantizer_mode == qcfg.QuantConfig._QuantizerMode.PT2E_DYNAMIC:
+      converter._experimental_qdq_conversion_mode = "DYNAMIC"
+    elif quantizer_mode == qcfg.QuantConfig._QuantizerMode.PT2E_STATIC:
+      converter._experimental_qdq_conversion_mode = "STATIC"

ai_edge_torch/{convert → _convert}/converter.py

@@ -15,20 +15,23 @@
 
 from __future__ import annotations
 
-from typing import Any, Dict, Optional, Tuple, Union
-
-import torch
+from typing import Any, Optional, Tuple, Union
 
 from ai_edge_torch import model
-from ai_edge_torch.convert import conversion
-from ai_edge_torch.convert import conversion_utils as cutils
+from ai_edge_torch._convert import conversion
+from ai_edge_torch._convert import signature as signature_module
 from ai_edge_torch.quantize import quant_config as qcfg
+import torch
 
 
 class Converter:
+  """A converter for converting PyTorch models to edge models.
+
+  This class allows adding multiple signatures to the converted edge model.
+  """
 
   def __init__(self):
-    self._signatures: list[cutils.Signature] = []
+    self._signatures: list[signature_module.Signature] = []
 
   def signature(
       self,
@@ -37,9 +40,9 @@ class Converter:
       sample_args=None,
       sample_kwargs=None,
       *,
-      dynamic_shapes: Optional[Union[Dict[str, Any], Tuple[Any]]] = None,
+      dynamic_shapes: Optional[Union[dict[str, Any], Tuple[Any, ...]]] = None,
   ) -> Converter:
-    """Alias to `add_signature`"""
+    """Functions as an alias to `add_signature`."""
     return self.add_signature(
         name, module, sample_args, sample_kwargs, dynamic_shapes=dynamic_shapes
     )
@@ -51,31 +54,44 @@ class Converter:
       sample_args=None,
       sample_kwargs=None,
       *,
-      dynamic_shapes: Optional[Union[Dict[str, Any], Tuple[Any]]] = None,
+      dynamic_shapes: Optional[Union[dict[str, Any], Tuple[Any, ...]]] = None,
   ) -> Converter:
     """Allows adding a new named torch model along with sample args to the conversion.
 
     Args:
       name: The name of the signature included in the converted edge model.
       module: The torch module to be converted.
-      sample_args: Tuple of tensors by which the torch module will be traced with prior to conversion.
-      sample_kwargs: Dict of str to tensor by which the torch module will be traced with prior to conversion.
-      dynamic_shapes: Optional dict or tuple that specify dynamic shape specifications for each input in original order.
-        See https://pytorch.org/docs/stable/export.html#expressing-dynamism for more details.
+      sample_args: Tuple of tensors by which the torch module will be traced
+        with prior to conversion.
+      sample_kwargs: Dict of str to tensor by which the torch module will be
+        traced with prior to conversion.
+      dynamic_shapes: Optional dict or tuple that specify dynamic shape
+        specifications for each input in original order. See
+        https://pytorch.org/docs/stable/export.html#expressing-dynamism for more
+          details.
+
+    Returns:
+      The converter object itself.
 
     Raises:
       ValueError: If a signature with the provided name already exists.
     """
 
     if name in [sig.name for sig in self._signatures]:
-      raise ValueError(f"A signature with the provided name ({name}) is already added.")
+      raise ValueError(
+          f"A signature with the provided name ({name}) is already added."
+      )
 
     if sample_args is None and sample_kwargs is None:
       raise ValueError("sample_args or sample_kwargs must be provided.")
 
     self._signatures.append(
-        cutils.Signature(
-            name, module, sample_args, sample_kwargs, dynamic_shapes=dynamic_shapes
+        signature_module.Signature(
+            name,
+            module,
+            sample_args,
+            sample_kwargs,
+            dynamic_shapes=dynamic_shapes,
         )
     )
     return self
@@ -87,8 +103,8 @@ class Converter:
       sample_kwargs=None,
       *,
       quant_config: Optional[qcfg.QuantConfig] = None,
-      dynamic_shapes: Optional[Union[Dict[str, Any], Tuple[Any]]] = None,
-      _ai_edge_converter_flags: dict = {},
+      dynamic_shapes: Optional[Union[dict[str, Any], Tuple[Any, ...]]] = None,
+      _ai_edge_converter_flags: Optional[dict[str, Any]] = None,
   ) -> model.TfLiteModel:
     """Finalizes the conversion and produces an edge model.
 
@@ -96,31 +112,44 @@ class Converter:
 
       edge_model = Converter().signature(name, module, args).convert()
 
-    Or it could be used to set the default signature for the converted edge model:
+    Or it could be used to set the default signature for the converted edge
+    model:
 
       edge_model =  Converter().convert(module, args)
 
     Args:
-      name: The name of the signature included in the converted edge model.
       module: The torch module to be converted.
-      sample_args: Tuple of tensors by which the torch module will be traced with prior to conversion.
-      sample_kwargs: Dict of str to tensor by which the torch module will be traced with prior to conversion.
+      sample_args: Tuple of tensors by which the torch module will be traced
+        with prior to conversion.
+      sample_kwargs: Dict of str to tensor by which the torch module will be
+        traced with prior to conversion.
       quant_config: User-defined quantization method and scheme of the model.
-      dynamic_shapes: Optional dict or tuple that specify dynamic shape specifications for each input in original order.
-        See https://pytorch.org/docs/stable/export.html#expressing-dynamism for more details.
-      _ai_edge_converter_flags: A nested dictionary allowing setting flags for the underlying converter.
-        This gives access to an implementation detail of this function and so needs to be treated as such.
-        Please do not rely on this parameter except for local debugging as this can be removed in a future release.
+      dynamic_shapes: Optional dict or tuple that specify dynamic shape
+        specifications for each input in original order. See
+        https://pytorch.org/docs/stable/export.html#expressing-dynamism for more
+          details.
+      _ai_edge_converter_flags: A nested dictionary allowing setting flags for
+        the underlying converter. This gives access to an implementation detail
+        of this function and so needs to be treated as such. Please do not rely
+        on this parameter except for local debugging as this can be removed in a
+        future release.
+
+    Returns:
+      The converted edge model.
 
     Raises:
-      ValueError: If the arguments are not provided as expected. See the example in this functions's comment.
+      ValueError: If the arguments are not provided as expected. See the example
+      in this functions's comment.
     """
+    if _ai_edge_converter_flags is None:
+      _ai_edge_converter_flags = {}
+
     if module is not None:
       if (
           sample_args is not None or sample_kwargs is not None
       ):  # both module and args provided
         self.add_signature(
-            cutils.DEFAULT_SIGNATURE_NAME,
+            model.DEFAULT_SIGNATURE_NAME,
             module,
             sample_args,
             sample_kwargs,
@@ -128,9 +157,9 @@ class Converter:
         )
       else:  # module is provided but not args
         raise ValueError(
-            "sample_args or sample_kwargs must be provided if a module is specified."
+            "sample_args or sample_kwargs must be provided if a module is"
+            " specified."
         )
-
     return conversion.convert_signatures(
         self._signatures,
         quant_config=quant_config,
@@ -143,22 +172,28 @@ def signature(
     module: torch.nn.Module,
     sample_args=None,
     sample_kwargs=None,
-    dynamic_shapes: Optional[Union[Dict[str, Any], Tuple[Any]]] = None,
+    dynamic_shapes: Optional[Union[dict[str, Any], Tuple[Any, ...]]] = None,
 ) -> Converter:
   """Initiates a Converter object with the provided signature.
 
   Args:
     name: The name of the signature included in the converted edge model.
     module: The torch module to be converted.
-    sample_args: Tuple of tensors by which the torch module will be traced with prior to conversion.
-    sample_kwargs: Dict of str to tensor by which the torch module will be traced with prior to conversion.
-    dynamic_shapes: Optional dict or tuple that specify dynamic shape specifications for each input in original order.
-      See https://pytorch.org/docs/stable/export.html#expressing-dynamism for more details.
+    sample_args: Tuple of tensors by which the torch module will be traced with
+      prior to conversion.
+    sample_kwargs: Dict of str to tensor by which the torch module will be
+      traced with prior to conversion.
+    dynamic_shapes: Optional dict or tuple that specify dynamic shape
+      specifications for each input in original order. See
+      https://pytorch.org/docs/stable/export.html#expressing-dynamism for more
+        details.
+
+  Returns:
+    A Converter object with the provided signature.
 
   Example:
     converter = ai_edge_torch.signature(name, module, args)
     edge_model = converter.convert()
-
   """
   return Converter().signature(
       name, module, sample_args, sample_kwargs, dynamic_shapes=dynamic_shapes
@@ -171,27 +206,38 @@ def convert(
     sample_kwargs=None,
     *,
     quant_config: Optional[qcfg.QuantConfig] = None,
-    dynamic_shapes: Optional[Union[Dict[str, Any], Tuple[Any]]] = None,
-    _ai_edge_converter_flags: dict = {},
+    dynamic_shapes: Optional[Union[dict[str, Any], Tuple[Any, ...]]] = None,
+    _ai_edge_converter_flags: Optional[dict[str, Any]] = None,
 ) -> model.TfLiteModel:
-  """Allows converting a PyTorch model to an edge model with one default signature in one step.
+  """Converts a PyTorch model to an edge model with a default signature.
 
   Args:
     module: The torch module to be converted.
-    sample_args: Tuple of tensors by which the torch module will be traced with prior to conversion.
-    sample_kwargs: Dict of str to tensor by which the torch module will be traced with prior to conversion.
+    sample_args: Tuple of tensors by which the torch module will be traced with
+      prior to conversion.
+    sample_kwargs: Dict of str to tensor by which the torch module will be
+      traced with prior to conversion.
     quant_config: User-defined quantization method and scheme of the model.
-    dynamic_shapes: Optional dict or tuple that specify dynamic shape specifications for each input in original order.
-      See https://pytorch.org/docs/stable/export.html#expressing-dynamism for more details.
-    _ai_edge_converter_flags: A nested dictionary allowing setting flags for the underlying converter.
-      This gives access to an implementation detail of this function and so needs to be treated as such.
-      Please do not rely on this parameter except for local debugging as this can be removed in a future release.
+    dynamic_shapes: Optional dict or tuple that specify dynamic shape
+      specifications for each input in original order. See
+      https://pytorch.org/docs/stable/export.html#expressing-dynamism for more
+        details.
+    _ai_edge_converter_flags: A nested dictionary allowing setting flags for the
+      underlying converter. This gives access to an implementation detail of
+      this function and so needs to be treated as such. Please do not rely on
+      this parameter except for local debugging as this can be removed in a
+      future release.
+
+  Returns:
+    The converted edge model.
 
   Example:
     edge_model = ai_edge_torch.convert(module, args)
-
   """
 
+  if _ai_edge_converter_flags is None:
+    _ai_edge_converter_flags = {}
+
   return Converter().convert(
       module,
       sample_args,

ai_edge_torch/_convert/fx_passes/__init__.py

@@ -0,0 +1,22 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from typing import Sequence, Union
+
+from ai_edge_torch._convert.fx_passes.build_aten_composite_pass import BuildAtenCompositePass
+from ai_edge_torch._convert.fx_passes.build_interpolate_composite_pass import BuildInterpolateCompositePass
+from ai_edge_torch._convert.fx_passes.inject_mlir_debuginfo_pass import InjectMlirDebuginfoPass
+from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass import OptimizeLayoutTransposesPass
+from ai_edge_torch.fx_pass_base import CanonicalizePass