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

ai_edge_torch/_convert/test/test_convert_composites.py

@@ -0,0 +1,234 @@
+# 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.
+# ==============================================================================
+"""Tests conversion modules that are meant to be wrapped as composites."""
+
+from collections.abc import Callable
+
+import ai_edge_torch
+from ai_edge_torch.testing import model_coverage
+import parameterized
+import torch
+
+from absl.testing import absltest as googletest
+
+
+def _func_to_torch_module(func: Callable[..., torch.Tensor]):
+  """Wraps a function into a torch module."""
+
+  class TestModule(torch.nn.Module):
+
+    def __init__(self, func):
+      super().__init__()
+      self._func = func
+
+    def forward(self, *args, **kwargs):
+      return self._func(*args, **kwargs)
+
+  return TestModule(func).eval()
+
+
+class TestConvertComposites(googletest.TestCase):
+  """Tests conversion modules that are meant to be wrapped as composites."""
+
+  def test_convert_hardswish(self):
+    """Tests conversion of a HardSwish module."""
+
+    args = (torch.randn((5, 10)),)
+    torch_module = torch.nn.Hardswish().eval()
+    edge_model = ai_edge_torch.convert(torch_module, args)
+
+    self.assertTrue(
+        model_coverage.compare_tflite_torch(edge_model, torch_module, args)
+    )
+
+  @parameterized.parameterized.expand([
+      # (input_size, kernel_size, stride, padding, ceil_mode,
+      # count_include_pad, divisor_override)
+      # no padding, stride = 1
+      ([1, 3, 6, 6], [3, 3], [1, 1], [0, 0], False, True, None),
+      # add stride
+      ([1, 3, 6, 6], [3, 3], [2, 2], [0, 0], False, True, None),
+      # default values
+      ([1, 3, 6, 6], [3, 3]),
+      # add padding
+      ([1, 3, 6, 6], [3, 3], [1, 1], [1, 1], False, True, None),
+      # add different padding for different dims
+      ([1, 3, 6, 6], [3, 3], [1, 1], [0, 1], False, True, None),
+      # add both stride and padding
+      ([1, 3, 6, 6], [3, 3], [2, 2], [1, 1], False, True, None),
+      # padding set to one number
+      ([1, 3, 6, 6], [3, 3], [1, 1], 1, False, True, None),
+      # count_include_pad = False
+      ([1, 3, 6, 6], [3, 3], [1, 1], [1, 1], False, False, None),
+      # ceil_mode = True
+      ([1, 3, 6, 6], [3, 3], [1, 1], [1, 1], True, True, None),
+      # ceil_mode = True, stride=[3, 3]
+      ([1, 3, 6, 6], [3, 3], [3, 3], [1, 1], True, True, None),
+      # set divisor_override
+      ([1, 3, 6, 6], [3, 3], [1, 1], 0, False, True, 6),
+  ])
+  def test_convert_avg_pool2d(self, input_size, *args):
+    """Tests conversion of a module containing an avg_pool2d aten."""
+    torch_module = _func_to_torch_module(
+        lambda input_tensor: torch.ops.aten.avg_pool2d(input_tensor, *args)
+    )
+    tracing_args = (torch.randn(*input_size),)
+    edge_model = ai_edge_torch.convert(torch_module, tracing_args)
+
+    self.assertTrue(
+        model_coverage.compare_tflite_torch(
+            edge_model, torch_module, tracing_args
+        )
+    )
+
+  @parameterized.parameterized.expand([
+      # use scale_factor with align_corners=False
+      (
+          [1, 3, 10, 10],
+          dict(scale_factor=3.0, mode='bilinear', align_corners=False),
+      ),
+      # use scale_factor with align_corners=true
+      (
+          [1, 3, 10, 10],
+          dict(scale_factor=3.0, mode='bilinear', align_corners=True),
+      ),
+      # use size
+      ([1, 3, 10, 10], dict(size=[15, 20], mode='bilinear')),
+      # use size with align_corners=true
+      (
+          [1, 3, 10, 10],
+          dict(size=[15, 20], mode='bilinear', align_corners=True),
+      ),
+  ])
+  def test_convert_upsample_bilinear_functional(self, input_size, kwargs):
+    """Tests conversion of a torch.nn.functional.upsample module."""
+    torch_module = _func_to_torch_module(
+        lambda input_tensor: torch.nn.functional.upsample(  # pylint: disable=unnecessary-lambda
+            input_tensor, **kwargs
+        )
+    )
+    tracing_args = (torch.randn(*input_size),)
+    edge_model = ai_edge_torch.convert(torch_module, tracing_args)
+
+    self.assertTrue(
+        model_coverage.compare_tflite_torch(
+            edge_model, torch_module, tracing_args
+        )
+    )
+
+  @parameterized.parameterized.expand([
+      # use scale_factor with align_corners=False
+      (
+          [1, 3, 10, 10],
+          dict(scale_factor=3.0, mode='bilinear', align_corners=False),
+      ),
+      # use scale_factor with align_corners=true
+      (
+          [1, 3, 10, 10],
+          dict(scale_factor=3.0, mode='bilinear', align_corners=True),
+      ),
+      # use size
+      ([1, 3, 10, 10], dict(size=[15, 20], mode='bilinear')),
+      # use size with align_corners=true
+      (
+          [1, 3, 10, 10],
+          dict(size=[15, 20], mode='bilinear', align_corners=True),
+      ),
+  ])
+  def test_convert_upsample_bilinear(self, input_size, kwargs):
+    """Tests conversion of a torch.nn.Upsample module."""
+    torch_module = _func_to_torch_module(
+        lambda input_tensor: torch.nn.Upsample(**kwargs)(input_tensor)  # pylint: disable=unnecessary-lambda
+    )
+    tracing_args = (torch.randn(*input_size),)
+    edge_model = ai_edge_torch.convert(torch_module, tracing_args)
+
+    self.assertTrue(
+        model_coverage.compare_tflite_torch(
+            edge_model, torch_module, tracing_args
+        )
+    )
+
+  @parameterized.parameterized.expand([
+      # use scale_factor with align_corners=False
+      (
+          [1, 3, 10, 10],
+          dict(scale_factor=3.0, mode='bilinear', align_corners=False),
+      ),
+      # use scale_factor with align_corners=true
+      (
+          [1, 3, 10, 10],
+          dict(scale_factor=3.0, mode='bilinear', align_corners=True),
+      ),
+      # use size
+      ([1, 3, 10, 10], dict(size=[15, 20], mode='bilinear')),
+      # use size with align_corners=true
+      (
+          [1, 3, 10, 10],
+          dict(size=[15, 20], mode='bilinear', align_corners=True),
+      ),
+  ])
+  def test_convert_interpolate_bilinear_functional(self, input_size, kwargs):
+    """Tests conversion of a torch.nn.functional.interpolate module."""
+    torch_module = _func_to_torch_module(
+        lambda input_tensor: torch.nn.functional.interpolate(  # pylint: disable=unnecessary-lambda
+            input_tensor, **kwargs
+        )
+    )
+    tracing_args = (torch.randn(*input_size),)
+    edge_model = ai_edge_torch.convert(torch_module, tracing_args)
+
+    self.assertTrue(
+        model_coverage.compare_tflite_torch(
+            edge_model, torch_module, tracing_args
+        )
+    )
+
+  def test_convert_gelu(self):
+    """Tests conversion of a GELU module."""
+
+    args = (torch.randn((5, 10)),)
+    torch_module = torch.nn.GELU().eval()
+    edge_model = ai_edge_torch.convert(torch_module, args)
+
+    self.assertTrue(
+        model_coverage.compare_tflite_torch(edge_model, torch_module, args)
+    )
+
+  def test_convert_gelu_approximate(self):
+    """Tests conversion of an Approximate GELU module."""
+
+    args = (torch.randn((5, 10)),)
+    torch_module = torch.nn.GELU('tanh').eval()
+    edge_model = ai_edge_torch.convert(torch_module, args)
+
+    self.assertTrue(
+        model_coverage.compare_tflite_torch(edge_model, torch_module, args)
+    )
+
+  def test_convert_embedding_lookup(self):
+    """Tests conversion of an Embedding module."""
+
+    args = (torch.full((1, 10), 0, dtype=torch.long),)
+    torch_module = torch.nn.Embedding(10, 10)
+    edge_model = ai_edge_torch.convert(torch_module, args)
+
+    self.assertTrue(
+        model_coverage.compare_tflite_torch(edge_model, torch_module, args)
+    )
+
+
+if __name__ == '__main__':
+  googletest.main()

ai_edge_torch/_convert/test/test_convert_multisig.py

@@ -0,0 +1,189 @@
+# 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.
+# ==============================================================================
+"""Tests for multi-signature conversion."""
+
+import ai_edge_torch
+from ai_edge_torch.testing import model_coverage
+import torch
+from torch import nn
+
+from absl.testing import absltest as googletest
+
+
+class FullyConnectedModel(nn.Module):
+  """A simple fully connected model with two fully connected layers."""
+
+  def __init__(self, input_size, hidden_size, output_size):
+    super(FullyConnectedModel, self).__init__()
+    self.fc = nn.Linear(input_size, hidden_size)  # Fully connected layer
+    self.relu = nn.ReLU()  # Activation function
+    self.output = nn.Linear(hidden_size, output_size)
+
+  def forward(self, x):
+    x = self.fc(x)
+    x = self.relu(x)
+    x = self.output(x)
+    return x
+
+
+class FullyConvModel(nn.Module):
+  """A simple fully convolutional model with two convolutions."""
+
+  def __init__(self):
+    super(FullyConvModel, self).__init__()
+    self.conv1 = nn.Conv2d(
+        3, 16, kernel_size=3, padding=1
+    )  # Input channels: 3 (RGB), Output channels: 16
+    self.relu = nn.ReLU(inplace=True)
+    self.conv2 = nn.Conv2d(
+        16, 1, kernel_size=1
+    )  # Output channels: 1 (single channel output)
+
+  def forward(self, x):
+    x = self.conv1(x)
+    x = self.relu(x)
+    x = self.conv2(x)
+    return x
+
+
+class TestConvertMultiSignature(googletest.TestCase):
+  """Tests conversion of various modules through multi-signature conversion."""
+
+  def setUp(self):
+    super().setUp()
+    torch.manual_seed(0)
+
+  def test_convert_with_default(self):
+    """Tests conversion of a model with two signatures one of which is the default."""
+    torch_module = FullyConvModel().eval()
+
+    args = (torch.randn(4, 3, 12, 12),)
+    large_args = (torch.randn(4, 3, 24, 24),)
+
+    signature_name = "large_input"
+
+    edge_model = ai_edge_torch.signature(
+        signature_name, torch_module, large_args
+    ).convert(torch_module, args)
+
+    self.assertTrue(
+        model_coverage.compare_tflite_torch(edge_model, torch_module, args)
+    )
+    self.assertTrue(
+        model_coverage.compare_tflite_torch(
+            edge_model, torch_module, large_args, signature_name=signature_name
+        )
+    )
+
+  def test_convert_no_default(self):
+    """Tests conversion of a model with two signatures none of which is the default."""
+    torch_module = FullyConvModel().eval()
+
+    args = (torch.randn(4, 3, 12, 12),)
+    large_args = (torch.randn(4, 3, 24, 24),)
+
+    signature_name_1 = "input"
+    signature_name_2 = "large_input"
+
+    edge_model = (
+        ai_edge_torch.signature(signature_name_1, torch_module, args)
+        .signature(signature_name_2, torch_module, large_args)
+        .convert()
+    )
+
+    with self.assertRaises(ValueError):
+      edge_model(*args)
+
+    self.assertTrue(
+        model_coverage.compare_tflite_torch(
+            edge_model, torch_module, args, signature_name=signature_name_1
+        )
+    )
+    self.assertTrue(
+        model_coverage.compare_tflite_torch(
+            edge_model,
+            torch_module,
+            large_args,
+            signature_name=signature_name_2,
+        )
+    )
+
+  def test_convert_signature_helper(self):
+    """Tests the ai_edge_torch.signature helper function works."""
+    torch_module = FullyConvModel().eval()
+
+    args = (torch.randn(4, 3, 12, 12),)
+    large_args = (torch.randn(4, 3, 24, 24),)
+
+    signature_name = "large_input"
+
+    edge_model = ai_edge_torch.signature(
+        signature_name, torch_module, large_args
+    ).convert(torch_module, args)
+
+    self.assertTrue(
+        model_coverage.compare_tflite_torch(edge_model, torch_module, args)
+    )
+    self.assertTrue(
+        model_coverage.compare_tflite_torch(
+            edge_model, torch_module, large_args, signature_name=signature_name
+        )
+    )
+
+  def test_convert_separate_modules(self):
+    """Tests conversion of two completely different modules as separate signatures."""
+    fully_conv = FullyConvModel().eval()
+    fully_connected = FullyConnectedModel(10, 5, 10).eval()
+
+    fully_conv_args = (torch.randn(4, 3, 12, 12),)
+    fully_connected_args = (torch.randn(10),)
+
+    fully_conv_signature_name = "fully_conv"
+    fully_connected_signature_name = "fully_connected"
+
+    edge_model = (
+        ai_edge_torch.signature(
+            fully_conv_signature_name, fully_conv, fully_conv_args
+        )
+        .signature(
+            fully_connected_signature_name,
+            fully_connected,
+            fully_connected_args,
+        )
+        .convert(fully_connected, fully_connected_args)
+    )
+
+    fully_conv_inference_args = (torch.randn(4, 3, 12, 12),)
+    fully_connected_inference_args = (torch.randn(10),)
+    self.assertTrue(
+        model_coverage.compare_tflite_torch(
+            edge_model,
+            fully_conv,
+            fully_conv_inference_args,
+            signature_name=fully_conv_signature_name,
+        )
+    )
+    self.assertTrue(
+        model_coverage.compare_tflite_torch(
+            edge_model,
+            fully_connected,
+            fully_connected_inference_args,
+            signature_name=fully_connected_signature_name,
+        )
+    )
+
+
+if __name__ == "__main__":
+  googletest.main()

ai_edge_torch/_convert/test/test_to_channel_last_io.py

@@ -0,0 +1,96 @@
+# 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.
+# ==============================================================================
+"""Tests for to_channel_last_io API and module wrapper."""
+
+import ai_edge_torch
+import torch
+
+from absl.testing import absltest as googletest
+
+
+class Identity(torch.nn.Module):
+
+  def forward(self, x):
+    return x
+
+
+class TestToChannelLastIO(googletest.TestCase):
+  """Tests to_channel_last_io API and module wrapper."""
+
+  def test_no_transformations(self):
+    x = torch.rand(1, 3, 10, 10)
+    y = ai_edge_torch.to_channel_last_io(Identity())(x)
+    self.assertEqual(y.shape, (1, 3, 10, 10))
+
+  def test_args(self):
+    x = torch.rand(1, 10, 10, 3)
+    y = ai_edge_torch.to_channel_last_io(Identity(), args=[0])(x)
+    self.assertEqual(y.shape, (1, 3, 10, 10))
+
+  def test_outputs(self):
+    x = torch.rand(1, 3, 10, 10)
+    y = ai_edge_torch.to_channel_last_io(Identity(), outputs=[0])(x)
+    self.assertEqual(y.shape, (1, 10, 10, 3))
+
+  def test_args_outputs(self):
+    x = torch.rand(1, 10, 10, 3)
+    y = ai_edge_torch.to_channel_last_io(Identity(), args=[0], outputs=[0])(x)
+    self.assertEqual(y.shape, (1, 10, 10, 3))
+
+  def test_args_5d(self):
+    x = torch.rand(1, 10, 10, 10, 3)
+    y = ai_edge_torch.to_channel_last_io(Identity(), args=[0])(x)
+    self.assertEqual(y.shape, (1, 3, 10, 10, 10))
+
+  def test_outputs_5d(self):
+    x = torch.rand(1, 3, 10, 10, 10)
+    y = ai_edge_torch.to_channel_last_io(Identity(), outputs=[0])(x)
+    self.assertEqual(y.shape, (1, 10, 10, 10, 3))
+
+  def test_chained_wrappers(self):
+    x = torch.rand(1, 10, 10, 3)
+
+    m = Identity()
+    m = ai_edge_torch.to_channel_last_io(m, args=[0])
+    m = ai_edge_torch.to_channel_last_io(m, outputs=[0])
+
+    y = m(x)
+    self.assertEqual(y.shape, (1, 10, 10, 3))
+
+  def test_list_args(self):
+    class Add(torch.nn.Module):
+
+      def forward(self, x, y):
+        return x + y
+
+    x = (torch.rand(1, 10, 10, 3), torch.rand(1, 10, 10, 3))
+    y = ai_edge_torch.to_channel_last_io(Add(), args=[0, 1])(*x)
+    self.assertEqual(y.shape, (1, 3, 10, 10))
+
+  def test_list_outputs(self):
+    class TwoIdentity(torch.nn.Module):
+
+      def forward(self, x):
+        return x, x
+
+    x = torch.rand(1, 3, 10, 10)
+    y = ai_edge_torch.to_channel_last_io(TwoIdentity(), outputs=[0])(x)
+    self.assertIsInstance(y, tuple)
+    self.assertEqual(y[0].shape, (1, 10, 10, 3))
+    self.assertEqual(y[1].shape, (1, 3, 10, 10))
+
+
+if __name__ == "__main__":
+  googletest.main()

ai_edge_torch/_convert/to_channel_last_io.py

@@ -0,0 +1,92 @@
+# 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.
+# ==============================================================================
+"""Transforms the input and output of a module to channel last layout."""
+
+from typing import Optional
+
+import torch
+from torch import nn
+
+
+class ChannelLastIOWrapper(nn.Module):
+
+  def __init__(self, wrapped, *, args=None, outputs=None):
+    super().__init__()
+    self.wrapped = wrapped
+    self._args = args or []
+    self._outputs = outputs or []
+
+  def _to_channel_last(self, x):
+    if not torch.is_tensor(x):
+      raise ValueError("Input must be a torch tensor")
+    if x.ndim < 3:
+      raise ValueError(
+          "Input must be a tensor with rank >= 3 in layout (N, C, ...)"
+      )
+    dims = [0, *range(2, x.ndim), 1]
+    return torch.permute(x, dims)
+
+  def _to_channel_first(self, x):
+    if not torch.is_tensor(x):
+      raise ValueError("Input must be a torch tensor.")
+    if x.ndim < 3:
+      raise ValueError(
+          "Input must be a tensor with rank >= 3 in layout (N, ..., C)"
+      )
+    dims = [0, x.ndim - 1, *range(1, x.ndim - 1)]
+    return torch.permute(x, dims)
+
+  def forward(self, *args, **kwargs):
+    args = list(args)
+    for i in self._args:
+      args[i] = self._to_channel_first(args[i])
+
+    outputs = self.wrapped(*args, **kwargs)
+
+    if not isinstance(outputs, (list, tuple)):
+      outputs_is_list = False
+      output_list = [outputs]
+    else:
+      outputs_is_list = True
+      output_list = list(outputs)
+
+    for i in self._outputs:
+      output_list[i] = self._to_channel_last(output_list[i])
+
+    if not outputs_is_list:
+      return output_list[0]
+    else:
+      return type(outputs)(output_list)
+
+
+def to_channel_last_io(
+    module: nn.Module,
+    args: Optional[list[int]] = None,
+    outputs: Optional[list[int]] = None,
+):
+  """Wraps the module with channel first to channel last layout transformations.
+
+  Args:
+    args (list[int]): Transform args with indices in the list from channel first
+      (N, C, ...) to channel last (N, ..., C).
+    outputs (list[int]): Transform outputs with indices in the list from channel
+      first (N, C, ...) to channel last (N, ..., C).
+
+  Returns:
+    The wrapped nn.Module with additional layout transposes after inputs and/or
+    before
+    outputs.
+  """
+  return ChannelLastIOWrapper(module, args=args, outputs=outputs)

ai_edge_torch/conftest.py

@@ -0,0 +1,20 @@
+# 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 absl import flags
+
+
+def pytest_configure(config):
+  flags.FLAGS.mark_as_parsed()

ai_edge_torch/debug/__init__.py

@@ -0,0 +1,17 @@
+# 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 .culprit import _search_model
+from .culprit import find_culprits