ai-edge-torch-nightly 0.2.0.dev20240714__py3-none-any.whl

ai_edge_torch/convert/test/test_convert_composites.py

@@ -0,0 +1,192 @@
+# 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 Callable
+import unittest
+
+import parameterized
+import torch
+
+import ai_edge_torch
+from ai_edge_torch.testing import model_coverage
+
+
+def _func_to_torch_module(func: Callable):
+  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(unittest.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),
+          # 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),
+          # padding set to one number
+          ([1, 3, 6, 6], [3, 3], [1, 1], 1, False, True, None),
+      ]
+  )
+  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(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)
+    )
+    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(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))
+
+
+if __name__ == '__main__':
+  unittest.main()

ai_edge_torch/convert/test/test_convert_multisig.py

@@ -0,0 +1,139 @@
+# 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 unittest
+
+import torch
+import torchvision
+
+import ai_edge_torch
+from ai_edge_torch.testing import model_coverage
+
+
+class TestConvertMultiSignature(unittest.TestCase):
+  """Tests conversion of various modules through multi-signature conversion."""
+
+  def setUp(self):
+    torch.manual_seed(0)
+
+  def test_convert_mobilenet_v2_with_default(self):
+    """Tests conversion of a model with two signatures one of which is the default."""
+    torch_module = torchvision.models.mobilenet_v2().eval()
+
+    args = (torch.randn(4, 3, 224, 224),)
+    large_args = (torch.randn(4, 3, 336, 336),)
+
+    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_mobilenet_v2_no_default(self):
+    """Tests conversion of a model with two signatures none of which is the default."""
+    torch_module = torchvision.models.mobilenet_v2().eval()
+
+    args = (torch.randn(4, 3, 224, 224),)
+    large_args = (torch.randn(4, 3, 336, 336),)
+
+    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_mobilenet_v2_signature_helper(self):
+    """Tests the ai_edge_torch.signature helper function works."""
+    torch_module = torchvision.models.mobilenet_v2().eval()
+
+    args = (torch.randn(4, 3, 224, 224),)
+    large_args = (torch.randn(4, 3, 336, 336),)
+
+    signature_name = "large_input"
+
+    edge_model = ai_edge_torch.signature(signature_name, torch_module, args).convert(
+        torch_module, large_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."""
+    mobilentv2 = torchvision.models.mobilenet_v2().eval()
+    resnet18 = torchvision.models.resnet18().eval()
+
+    mobilenet_args = (torch.randn(4, 3, 224, 224),)
+    resnet_args = (torch.randn(4, 3, 224, 224),)
+
+    mobilenet_signature_name = "mobilentv2"
+    resnet_signature_name = "resnet18"
+
+    edge_model = (
+        ai_edge_torch.signature(mobilenet_signature_name, mobilentv2, mobilenet_args)
+        .signature(resnet_signature_name, resnet18, resnet_args)
+        .convert(resnet18, resnet_args)
+    )
+
+    mobilenet_inference_args = (torch.randn(4, 3, 224, 224),)
+    resnet_inference_args = (torch.randn(4, 3, 224, 224),)
+    self.assertTrue(
+        model_coverage.compare_tflite_torch(
+            edge_model,
+            mobilentv2,
+            mobilenet_inference_args,
+            signature_name=mobilenet_signature_name,
+        )
+    )
+    self.assertTrue(
+        model_coverage.compare_tflite_torch(
+            edge_model,
+            resnet18,
+            resnet_inference_args,
+            signature_name=resnet_signature_name,
+        )
+    )
+
+
+if __name__ == "__main__":
+  unittest.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.
+# ==============================================================================
+
+import unittest
+
+import torch
+
+import ai_edge_torch
+
+
+class Identity(torch.nn.Module):
+
+  def forward(self, x):
+    return x
+
+
+class TestToChannelLastIO(unittest.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__":
+  unittest.main()

ai_edge_torch/convert/to_channel_last_io.py

@@ -0,0 +1,85 @@
+# 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 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/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