ai-edge-torch-nightly 0.2.0.dev20240714__py3-none-any.whl → 0.3.0.dev20240926__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 → _convert}/test/test_to_channel_last_io.py

@@ -12,12 +12,12 @@
 # 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 unittest
-
+import ai_edge_torch
 import torch
 
-import ai_edge_torch
+from absl.testing import absltest as googletest
 
 
 class Identity(torch.nn.Module):
@@ -26,7 +26,7 @@ class Identity(torch.nn.Module):
     return x
 
 
-class TestToChannelLastIO(unittest.TestCase):
+class TestToChannelLastIO(googletest.TestCase):
   """Tests to_channel_last_io API and module wrapper."""
 
   def test_no_transformations(self):
@@ -93,4 +93,4 @@ class TestToChannelLastIO(unittest.TestCase):
 
 
 if __name__ == "__main__":
-  unittest.main()
+  googletest.main()

ai_edge_torch/{convert → _convert}/to_channel_last_io.py

@@ -12,6 +12,7 @@
 # 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
 
@@ -31,7 +32,9 @@ class ChannelLastIOWrapper(nn.Module):
     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, ...)")
+      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)
 
@@ -39,7 +42,9 @@ class ChannelLastIOWrapper(nn.Module):
     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)")
+      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)
 
@@ -78,8 +83,10 @@ def to_channel_last_io(
       (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
+    The wrapped nn.Module with additional layout transposes after inputs and/or
+    before
     outputs.
   """
   return ChannelLastIOWrapper(module, args=args, outputs=outputs)

ai_edge_torch/config.py

@@ -0,0 +1,27 @@
+# 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 library."""
+
+import dataclasses
+import os
+
+
+@dataclasses.dataclass
+class Config:
+  use_torch_xla: bool = os.environ.get("USE_TORCH_XLA", "true").lower() in (
+      "1",
+      "true",
+  )

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()