ai-edge-torch-nightly 0.3.0.dev20250114__py3-none-any.whl

ai_edge_torch/__init__.py

@@ -0,0 +1,32 @@
+# 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 ai_edge_torch._config import config
+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:
+  """Imports an ai_edge_torch model from disk.
+
+  Args:
+    path: The path to the serialized ai_edge_torch model.
+
+  Returns:
+    An ai_edge_torch.model.Model object.
+  """
+  return Model.load(path)

ai_edge_torch/_config.py

@@ -0,0 +1,69 @@
+# 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.
+# ==============================================================================
+
+"""Provides a configuration for the ai-edge-torch."""
+
+import functools
+import logging
+import os
+
+__all__ = ["config"]
+
+
+def _get_bool_env_var(name: str, default: bool) -> bool:
+  var = os.environ.get(name, "false")
+  var = var.lower().strip()
+  if var in ("y", "yes", "t", "true", "on", "1"):
+    return True
+  elif var in ("n", "no", "f", "false", "off", "0"):
+    return False
+  else:
+    logging.warning("Invalid %s value is ignored: %s.", name, var)
+    return default
+
+
+class _Config:
+  """ai-edge-torch global configs."""
+
+  @property
+  @functools.cache  # pylint: disable=method-cache-max-size-none
+  def use_torch_xla(self) -> bool:
+    """True if using torch_xla to lower torch ops to StableHLO.
+
+    To use torch_xla as the lowering backend, set environment variable
+    `USE_TORCH_XLA` to "true".
+    """
+    return _get_bool_env_var("USE_TORCH_XLA", default=False)
+
+  @property
+  def in_oss(self) -> bool:
+    """True if the code is not running in google internal environment."""
+    return True
+
+  @property
+  def enable_group_norm_composite(self) -> bool:
+    """True if lowering group norm in StableHLO composite.
+
+    Currently only supports NHWC group norm generated by
+    OptimizeLayoutTransposesPass.
+    """
+    return _get_bool_env_var("ENABLE_GROUP_NORM_COMPOSITE", default=False)
+
+  @enable_group_norm_composite.setter
+  def enable_group_norm_composite(self, value: bool):
+    os.environ["ENABLE_GROUP_NORM_COMPOSITE"] = "y" if value else "n"
+
+
+config = _Config()

ai_edge_torch/_convert/__init__.py

@@ -0,0 +1,14 @@
+# 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.
+# ==============================================================================

ai_edge_torch/_convert/conversion.py

@@ -0,0 +1,153 @@
+# 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
+from typing import Any, Literal, Optional, Union
+
+import ai_edge_torch
+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
+
+
+def _run_convert_passes(
+    exported_program: torch.export.ExportedProgram,
+) -> torch.export.ExportedProgram:
+  exported_program = generative_fx_passes.run_generative_passes(
+      exported_program
+  )
+
+  passes = [
+      fx_passes.BuildInterpolateCompositePass(),
+      fx_passes.CanonicalizePass(),
+      fx_passes.OptimizeLayoutTransposesPass(),
+      fx_passes.CanonicalizePass(),
+      fx_passes.BuildAtenCompositePass(),
+      fx_passes.CanonicalizePass(),
+      fx_passes.RemoveNonUserOutputsPass(),
+      fx_passes.CanonicalizePass(),
+  ]
+
+  # Debuginfo is not injected automatically by odml_torch. Only inject
+  # debuginfo via fx pass when using torch_xla.
+  if ai_edge_torch.config.use_torch_xla:
+    passes += [
+        fx_passes.InjectMlirDebuginfoPass(),
+        fx_passes.CanonicalizePass(),
+    ]
+
+  exported_program = fx_pass_base.run_passes(exported_program, passes)
+  return exported_program
+
+
+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],
+    *,
+    strict_export: Union[Literal["auto"], bool] = True,
+    quant_config: Optional[qcfg.QuantConfig] = None,
+    _tfl_converter_flags: Optional[dict[str, Any]] = None,
+    _saved_model_dir: Optional[str] = None,
+) -> 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.
+      strict_export: Experimental `strict` arg for torch.export.export. When
+        enabled, the export function will trace the program through TorchDynamo
+        and ensure the soundness of the exported graph. When
+        strict_export="auto", the function will try to export module in both
+        modes and use the first one succeeds for downstream conversion.
+      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.
+      _saved_model_dir: Directory for the intermediate saved model. If not
+        specified, a random temporary directory would be used.
+
+  Returns:
+    The converted `model.TfLiteModel` object.
+  """
+  if _tfl_converter_flags is None:
+    _tfl_converter_flags = {}
+
+  _warn_training_modules(signatures)
+
+  def export(*args, **kwargs):
+    nonlocal strict_export
+    if strict_export == "auto":
+      try:
+        exported_program = torch.export.export(*args, **kwargs, strict=True)
+      except Exception:
+        logging.warning(
+            "torch.export.export(..., strict=True) failed. Retrying with"
+            " strict=False"
+        )
+        exported_program = torch.export.export(*args, **kwargs, strict=False)
+    elif not strict_export:
+      exported_program = torch.export.export(*args, **kwargs, strict=False)
+    else:
+      exported_program = torch.export.export(*args, **kwargs, strict=True)
+
+    if hasattr(torch._decomp, "_decomp_table_to_post_autograd_aten"):
+      # Available after torch 2.5.0: `_decomp_table_to_post_autograd_aten` is a
+      # stop-gap table which replicates the old behaviour of post-dispatch IR.
+      # This could help ensure the collection of aten ops remaining still as the
+      # implementation of torch.export changes.
+      exported_program = exported_program.run_decompositions(
+          torch._decomp._decomp_table_to_post_autograd_aten()
+      )
+    return exported_program
+
+  exported_programs: torch.export.ExportedProgram = [
+      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,
+      _saved_model_dir=_saved_model_dir,
+  )
+
+  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/converter.py

@@ -0,0 +1,270 @@
+# 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 __future__ import annotations
+
+from typing import Any, Literal, Optional, Tuple, Union
+
+from ai_edge_torch import model
+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[signature_module.Signature] = []
+
+  def signature(
+      self,
+      name: str,
+      module: torch.nn.Module,
+      sample_args=None,
+      sample_kwargs=None,
+      *,
+      dynamic_shapes: Optional[Union[dict[str, Any], Tuple[Any, ...]]] = None,
+  ) -> Converter:
+    """Functions as an alias to `add_signature`."""
+    return self.add_signature(
+        name, module, sample_args, sample_kwargs, dynamic_shapes=dynamic_shapes
+    )
+
+  def add_signature(
+      self,
+      name: str,
+      module: torch.nn.Module,
+      sample_args=None,
+      sample_kwargs=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.
+
+    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."
+      )
+
+    if sample_args is None and sample_kwargs is None:
+      raise ValueError("sample_args or sample_kwargs must be provided.")
+
+    self._signatures.append(
+        signature_module.Signature(
+            name,
+            module,
+            sample_args,
+            sample_kwargs,
+            dynamic_shapes=dynamic_shapes,
+        )
+    )
+    return self
+
+  def convert(
+      self,
+      module: torch.nn.Module = None,
+      sample_args=None,
+      sample_kwargs=None,
+      *,
+      strict_export: Union[Literal["auto"], bool] = True,
+      quant_config: Optional[qcfg.QuantConfig] = None,
+      dynamic_shapes: Optional[Union[dict[str, Any], Tuple[Any, ...]]] = None,
+      _ai_edge_converter_flags: Optional[dict[str, Any]] = None,
+      _saved_model_dir: Optional[str] = None,
+  ) -> model.TfLiteModel:
+    """Finalizes the conversion and produces an edge model.
+
+    This could be called with no arguments as follows:
+
+      edge_model = Converter().signature(name, module, args).convert()
+
+    Or it could be used to set the default signature for the converted edge
+    model:
+
+      edge_model =  Converter().convert(module, args)
+
+    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.
+      strict_export: Experimental `strict` arg for torch.export.export. When
+        enabled, the export function will trace the program through TorchDynamo
+        and ensure the soundness of the exported graph. When
+        strict_export="auto", the function will try to export module in both
+        modes and use the first one succeeds for downstream 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.
+      _saved_model_dir: Directory for the intermediate saved model. If not
+        specified, a random temporary directory would be used.
+
+    Returns:
+      The converted edge model.
+
+    Raises:
+      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(
+            model.DEFAULT_SIGNATURE_NAME,
+            module,
+            sample_args,
+            sample_kwargs,
+            dynamic_shapes=dynamic_shapes,
+        )
+      else:  # module is provided but not args
+        raise ValueError(
+            "sample_args or sample_kwargs must be provided if a module is"
+            " specified."
+        )
+    return conversion.convert_signatures(
+        self._signatures,
+        strict_export=strict_export,
+        quant_config=quant_config,
+        _tfl_converter_flags=_ai_edge_converter_flags,
+        _saved_model_dir=_saved_model_dir,
+    )
+
+
+def signature(
+    name: str,
+    module: torch.nn.Module,
+    sample_args=None,
+    sample_kwargs=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.
+
+  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
+  )
+
+
+def convert(
+    module: torch.nn.Module = None,
+    sample_args=None,
+    sample_kwargs=None,
+    *,
+    strict_export: Union[Literal["auto"], bool] = True,
+    quant_config: Optional[qcfg.QuantConfig] = None,
+    dynamic_shapes: Optional[Union[dict[str, Any], Tuple[Any, ...]]] = None,
+    _ai_edge_converter_flags: Optional[dict[str, Any]] = None,
+    _saved_model_dir: Optional[str] = None,
+) -> model.TfLiteModel:
+  """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.
+    strict_export: Experimental `strict` arg for torch.export.export. When
+      enabled, the export function will trace the program through TorchDynamo
+      and ensure the soundness of the exported graph. When strict_export="auto",
+      the function will try to export module in both modes and use the first one
+      succeeds for downstream 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.
+    _saved_model_dir: Directory for the intermediate saved model. If not
+      specified, a random temporary directory would be used.
+
+  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,
+      sample_kwargs,
+      strict_export=strict_export,
+      quant_config=quant_config,
+      dynamic_shapes=dynamic_shapes,
+      _ai_edge_converter_flags=_ai_edge_converter_flags,
+      _saved_model_dir=_saved_model_dir,
+  )

ai_edge_torch/_convert/fx_passes/__init__.py

@@ -0,0 +1,23 @@
+# 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._convert.fx_passes.remove_non_user_outputs_pass import RemoveNonUserOutputsPass
+from ai_edge_torch.fx_pass_base import CanonicalizePass