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

ai_edge_torch/convert/fx_passes/optimize_layout_transposes_pass/pass_body.py

@@ -0,0 +1,293 @@
+# 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 os
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.ao.quantization.quantize_pt2e
+from torch.export import ExportedProgram
+from torch.fx import GraphModule
+from torch.fx import Node
+import torch.utils._pytree as pytree
+
+from ai_edge_torch.convert.fx_passes import ExportedProgramPassBase
+from ai_edge_torch.convert.fx_passes import ExportedProgramPassResult
+from ai_edge_torch.convert.fx_passes.optimize_layout_transposes_pass import layout_check  # NOQA
+from ai_edge_torch.convert.fx_passes.optimize_layout_transposes_pass import layout_mark  # NOQA
+from ai_edge_torch.convert.fx_passes.optimize_layout_transposes_pass import layout_partitioners  # NOQA
+from ai_edge_torch.convert.fx_passes.optimize_layout_transposes_pass import layout_rewrite  # NOQA
+from ai_edge_torch.convert.fx_passes.optimize_layout_transposes_pass import utils  # NOQA
+
+TransposeFunc = Union[utils.tensor_to_nchw, utils.tensor_to_nhwc]
+
+
+class OptimizeLayoutTransposesPass(ExportedProgramPassBase):
+
+  def get_source_meta(self, node: torch.fx.Node):
+    keys = ["stack_trace", "nn_module_stack", "source_fn_stack", "from_node"]
+    meta = {}
+    for key in keys:
+      if key in node.meta:
+        meta[key] = node.meta[key]
+    return meta
+
+  def insert_t_q_dq(
+      self,
+      graph: torch.fx.Graph,
+      input_dq: torch.fx.Node,
+      target: torch.fx.Node,
+      transpose_func: TransposeFunc,
+      transpose_node_meta: dict,
+  ) -> list[torch.fx.Node]:
+    """
+    original:
+        input_dq -> target
+    insert the node as:
+        input_dq -> (T q dq) -> target
+    """
+    assert utils.is_dq_node(input_dq)
+
+    q_args = input_dq.args[1:]
+    q_kwargs = input_dq.kwargs
+    q_op, dq_op = utils.get_paired_q_dq_ops(input_dq.target)
+    with graph.inserting_before(target):
+      t = graph.call_function(transpose_func, (input_dq,))
+      # Q and DQ inserted here may required updating the `axis` arg when they
+      # are per_channel ops. However, instead of updating here, the nodes would
+      # be marked as NHWC/NCHW and applied rewriters after partitioning.
+      q = graph.call_function(q_op, (t,) + q_args, q_kwargs)
+      dq = graph.call_function(dq_op, (q,) + q_args, q_kwargs)
+
+    input_dq.meta = transpose_node_meta
+    t.meta = transpose_node_meta
+    q.meta = transpose_node_meta
+    dq.meta = self.get_source_meta(target)
+
+    target.replace_input_with(input_dq, dq)
+    return [t, q, dq]
+
+  def insert_dq_t_q(
+      self,
+      graph: torch.fx.Graph,
+      input_q: torch.fx.Node,
+      target: torch.fx.Node,
+      transpose_func: TransposeFunc,
+      transpose_node_meta: dict,
+  ) -> list[torch.fx.Node]:
+    """
+    original:
+        input_q -> target
+    insert the node as:
+        input_q -> (dq T q) -> target
+    """
+    assert utils.is_q_node(input_q)
+
+    q_args = input_q.args[1:]
+    q_kwargs = input_q.kwargs
+    q_op, dq_op = self.get_paired_q_dq_ops(input_q.target)
+    with graph.inserting_before(target):
+      # Q and DQ inserted here may required updating the `axis` arg when they
+      # are per_channel ops. However, instead of updating here, the nodes would
+      # be marked as NHWC/NCHW and applied rewriters after partitioning.
+      dq = graph.call_function(dq_op, (input_q,) + q_args, q_kwargs)
+      t = graph.call_function(transpose_func, (dq,))
+      q = graph.call_function(q_op, (t,) + q_args, q_kwargs)
+
+    dq.meta = transpose_node_meta
+    t.meta = transpose_node_meta
+    q.meta = transpose_node_meta
+
+    target.replace_input_with(input_q, q)
+    return [dq, t, q]
+
+  def insert_layout_transpose(
+      self,
+      graph: torch.fx.Graph,
+      input_node: torch.fx.Node,
+      target_node: torch.fx.Node,
+      transpose_func: TransposeFunc,
+      transpose_node_meta: dict,
+  ) -> None:
+    assert transpose_func in (utils.tensor_to_nchw, utils.tensor_to_nhwc)
+
+    # new_nodes only contains Q/DQ/Transpose nodes, which are all SISO.
+    # Insertion order input nodes -> output nodes
+    new_nodes = []
+
+    # Constraint Q2: the NHWC partition's entry and exit must not be output
+    # edges of Q/DQ ops that are connected to a constant/weight tensor.
+    while layout_mark.is_const_node(input_node) and (
+        utils.is_dq_node(input_node) or utils.is_q_node(input_node)
+    ):
+      with graph.inserting_before(target_node):
+        new_input_node = graph.node_copy(input_node)
+
+      target_node.replace_input_with(input_node, new_input_node)
+
+      new_nodes = [new_input_node] + new_nodes
+      input_node, target_node = new_input_node.args[0], new_input_node
+
+    if utils.is_q_node(input_node):
+      # Constraint Q3: when the entry and exit is right after a q op (occur after a (dq-op-q)
+      # triplet), the transpose must be added as a quantized transpose in (dq-T-q)
+      # input_q -> (dq T q) -> target
+      new_nodes = (
+          self.insert_dq_t_q(
+              graph,
+              input_node,
+              target_node,
+              transpose_func,
+              transpose_node_meta,
+          )
+          + new_nodes
+      )
+    elif utils.is_dq_node(input_node):
+      # Constraint Q1: the NHWC partition's entry and exit cannot be edges
+      # within (dq-op-q) triplet.
+      # input_dq -> (T q dq) -> target
+      new_nodes = (
+          self.insert_t_q_dq(
+              graph,
+              input_node,
+              target_node,
+              transpose_func,
+              transpose_node_meta,
+          )
+          + new_nodes
+      )
+    else:
+      # input -> target
+      with graph.inserting_before(target_node):
+        t = graph.call_function(transpose_func, (input_node,))
+      t.meta = transpose_node_meta
+      target_node.replace_input_with(input_node, t)
+      new_nodes = [t] + new_nodes
+
+    # Mark new nodes as NCHW or NHWC
+    # For all nodes before the transpose, mark it as input_marker
+    # For all nodes after the transpose (incl. transpose), mark it as output_marker
+    if transpose_func == utils.tensor_to_nchw:
+      input_marker, target_marker = (
+          layout_mark.mark_as_nhwc_node,
+          layout_mark.mark_as_nchw_node,
+      )
+    else:
+      input_marker, target_marker = (
+          layout_mark.mark_as_nchw_node,
+          layout_mark.mark_as_nhwc_node,
+      )
+
+    marker = input_marker
+    for node in new_nodes:
+      if node.target == transpose_func:
+        marker = target_marker
+      marker(node)
+    assert marker == target_marker
+
+  def input_to_nhwc(
+      self,
+      graph: torch.fx.Graph,
+      input_node: torch.fx.Node,
+      target_node: torch.fx.Node,
+  ) -> None:
+    if layout_mark.is_nhwc_node(input_node):
+      return
+
+    if not layout_check.is_4d(input_node):
+      raise AssertionError(
+          f"Attempting to convert non-NHWC compatible node to NHWC: {input_node}"
+      )
+
+    # Assign target node's source meta to the to_NHWC node, because the transpose
+    # is added for the existence of target node.
+    self.insert_layout_transpose(
+        graph,
+        input_node,
+        target_node,
+        utils.tensor_to_nhwc,
+        self.get_source_meta(target_node),
+    )
+
+  def input_to_nchw(
+      self,
+      graph: torch.fx.Graph,
+      input_node: torch.fx.Node,
+      target_node: torch.fx.Node,
+  ) -> None:
+    if layout_mark.is_nchw_node(input_node):
+      return
+
+    self.insert_layout_transpose(
+        graph,
+        input_node,
+        target_node,
+        utils.tensor_to_nchw,
+        self.get_source_meta(input_node),
+    )
+
+  def mark_const_nodes(self, exported_program: torch.export.ExportedProgram):
+    graph_module = exported_program.graph_module
+    graph = graph_module.graph
+
+    input_specs = exported_program.graph_signature.input_specs
+    non_user_input_names = set()
+    for spec in input_specs:
+      if spec.kind != torch.export.graph_signature.InputKind.USER_INPUT:
+        non_user_input_names.add(spec.arg.name)
+
+    for node in graph.nodes:
+      has_input_nodes = len(node.all_input_nodes) > 0
+      all_inputs_are_const = all(map(layout_mark.is_const_node, node.all_input_nodes))
+      if (
+          node.name in non_user_input_names
+          or (has_input_nodes and all_inputs_are_const)
+          or (node.op != "placeholder" and not has_input_nodes)
+      ):
+        layout_mark.mark_as_const_node(node)
+
+  def call(self, exported_program: torch.export.ExportedProgram):
+    self.mark_const_nodes(exported_program)
+
+    graph_module = exported_program.graph_module
+    partitioner = os.environ.get("AIEDGETORCH_LAYOUT_OPTIMIZE_PARTITIONER", None)
+    if partitioner == "MINCUT":
+      graph_module = layout_partitioners.min_cut.partition(graph_module)
+    elif partitioner == "GREEDY":
+      graph_module = layout_partitioners.greedy.partition(graph_module)
+    else:
+      # By default use min cut partitioner if possible
+      if layout_partitioners.min_cut.can_partition(graph_module):
+        graph_module = layout_partitioners.min_cut.partition(graph_module)
+      else:
+        graph_module = layout_partitioners.greedy.partition(graph_module)
+
+    graph = graph_module.graph
+    for node in list(graph.nodes):
+      if layout_mark.is_nhwc_node(node):
+        for input_node in layout_check.get_layout_sensitive_inputs(node):
+          self.input_to_nhwc(graph, input_node, node)
+        layout_rewrite.rewrite_nhwc_node(node)
+      else:
+        for input_node in layout_check.get_layout_sensitive_inputs(node):
+          # Note: for non-4D tensors input_to_nchw is always noop.
+          self.input_to_nchw(graph, input_node, node)
+
+    graph_module.graph.eliminate_dead_code()
+    graph_module.recompile()
+    graph_module.graph.lint()
+    # Mark const node again for debugging
+    self.mark_const_nodes(exported_program)
+
+    return ExportedProgramPassResult(exported_program, True)

ai_edge_torch/convert/fx_passes/optimize_layout_transposes_pass/utils.py

@@ -0,0 +1,62 @@
+# 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 torch
+import torch.ao.quantization.quantize_pt2e
+
+
+def tensor_to_nhwc(t: torch.Tensor):
+  return torch.ops.aten.permute(t.contiguous(), [0, 2, 3, 1]).contiguous()
+
+
+def tensor_to_nchw(t: torch.Tensor):
+  return torch.ops.aten.permute(t.contiguous(), [0, 3, 1, 2]).contiguous()
+
+
+def flatten_torch_op_overloads(op):
+  if isinstance(op, torch._ops.OpOverloadPacket):
+    return [getattr(op, overload) for overload in op.overloads()]
+  return [op]
+
+
+_TORCH_Q_OPS = [
+    torch.ops.quantized_decomposed.quantize_per_tensor.default,
+    torch.ops.quantized_decomposed.quantize_per_tensor.tensor,
+    torch.ops.quantized_decomposed.quantize_per_tensor.tensor2,
+    torch.ops.quantized_decomposed.quantize_per_channel.default,
+]
+
+_TORCH_DQ_OPS = [
+    torch.ops.quantized_decomposed.dequantize_per_tensor.default,
+    torch.ops.quantized_decomposed.dequantize_per_tensor.tensor,
+    torch.ops.quantized_decomposed.dequantize_per_tensor.tensor2,
+    torch.ops.quantized_decomposed.dequantize_per_channel.default,
+]
+
+
+def is_q_node(node: torch.fx.Node):
+  return node.target in _TORCH_Q_OPS
+
+
+def is_dq_node(node: torch.fx.Node):
+  return node.target in _TORCH_DQ_OPS
+
+
+def get_paired_q_dq_ops(op: Callable) -> tuple[Callable, Callable]:
+  for q, dq in zip(_TORCH_Q_OPS, _TORCH_DQ_OPS):
+    if op in (q, dq):
+      return q, dq
+  raise AssertionError(f"{op} is not a Q/DQ op.")

ai_edge_torch/convert/test/__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/test/test_convert.py

@@ -0,0 +1,311 @@
+# 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 os
+import tempfile
+import unittest
+
+import torch
+import torchvision
+
+import ai_edge_torch
+from ai_edge_torch.convert import conversion_utils as cutils
+from ai_edge_torch.testing import model_coverage
+
+
+class TestConvert(unittest.TestCase):
+  """Tests conversion of various modules."""
+
+  def setUp(self):
+    torch.manual_seed(0)
+
+  def test_convert_add(self):
+    """Tests conversion of a simple Add module."""
+
+    class Add(torch.nn.Module):
+
+      def forward(self, a, b):
+        return a + b
+
+    args = (
+        torch.randn((5, 10)),
+        torch.randn((5, 10)),
+    )
+    torch_module = Add().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_dot_add(self):
+    class DotAdd(torch.nn.Module):
+      """Tests conversion of a matrix multiplication followed by an add."""
+
+      def forward(self, a, b, c):
+        return a @ b + c
+
+    args = (
+        torch.randn((5, 10)),
+        torch.randn((10, 5)),
+        torch.randn((5, 5)),
+    )
+    torch_module = DotAdd().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_resnet18(self):
+    args = (torch.randn(4, 3, 224, 224),)
+    torch_module = torchvision.models.resnet18().eval()
+    edge_model = ai_edge_torch.convert(torch_module, args)
+
+    self.assertTrue(model_coverage.compare_tflite_torch(edge_model, torch_module, args))
+
+  def test_signature_args_ordering(self):
+    """Tests conversion of a model with more than 10 arguments."""
+
+    class AddChainWith11Args(torch.nn.Module):
+
+      def forward(
+          self,
+          arg0: "f32[64]",
+          arg1: "f32[64]",
+          arg2: "f32[64]",
+          arg3: "f32[64]",
+          arg4: "f32[64]",
+          arg5: "f32[64]",
+          arg6: "f32[64]",
+          arg7: "f32[64]",
+          arg8: "f32[64]",
+          arg9: "f32[64]",
+          arg10: "f32[64]",
+      ):
+        add0 = torch.add(arg0, arg1)
+        add1 = torch.add(add0, arg2)
+        add2 = torch.add(add1, arg3)
+        add3 = torch.add(add2, arg4)
+        add4 = torch.add(add3, arg5)
+        add5 = torch.add(add4, arg6)
+        add6 = torch.add(add5, arg7)
+        add7 = torch.add(add6, arg8)
+        add8 = torch.add(add7, arg9)
+        add9 = torch.add(add8, arg10)
+        return add9
+
+    sample_input = lambda: (
+        torch.rand((64,), dtype=torch.float32),
+        torch.rand((64,), dtype=torch.float32),
+        torch.rand((64,), dtype=torch.float32),
+        torch.rand((64,), dtype=torch.float32),
+        torch.rand((64,), dtype=torch.float32),
+        torch.rand((64,), dtype=torch.float32),
+        torch.rand((64,), dtype=torch.float32),
+        torch.rand((64,), dtype=torch.float32),
+        torch.rand((64,), dtype=torch.float32),
+        torch.rand((64,), dtype=torch.float32),
+        torch.rand((64,), dtype=torch.float32),
+    )
+    torch_model = AddChainWith11Args().eval()
+    edge_model = ai_edge_torch.convert(torch_model, sample_input())
+
+    result = model_coverage.compare_tflite_torch(
+        edge_model, torch_model, sample_input, num_valid_inputs=10
+    )
+    self.assertTrue(result)
+
+  def test_multi_output_model(self):
+    """Tests conversion of a model that returns multiple outputs."""
+
+    class BasicAddModelWithMultipleOutputs(torch.nn.Module):
+
+      def forward(self, arg0, arg1):
+        add0 = arg0 + arg1
+        mul0 = arg0 * arg1
+        return add0, mul0
+
+    sample_input = (
+        torch.rand((64,), dtype=torch.float32),
+        torch.rand((64,), dtype=torch.float32),
+    )
+
+    torch_model = BasicAddModelWithMultipleOutputs().eval()
+    edge_model = ai_edge_torch.convert(torch_model, sample_input)
+
+    result = model_coverage.compare_tflite_torch(edge_model, torch_model, sample_input)
+    self.assertTrue(result)
+
+  def test_12_outputs_model(self):
+    """Tests conversion of a model that returns multiple outputs."""
+
+    class BasicAddModelWithMultipleOutputs(torch.nn.Module):
+
+      def forward(self, arg0, arg1):
+        add0 = arg0 + arg1
+        mul0 = arg0 * arg1
+        add1 = add0 + mul0
+        mul1 = add0 * mul0
+        add2 = add1 + mul1
+        mul2 = add1 * mul1
+        add3 = add2 + mul2
+        mul3 = add2 * mul2
+        add4 = add3 + mul3
+        mul4 = add3 * mul3
+        add5 = add4 + mul4
+        mul5 = add4 * mul4
+
+        return (
+            add0,
+            mul0,
+            add1,
+            mul1,
+            add2,
+            mul2,
+            add3,
+            mul3,
+            add4,
+            mul4,
+            add5,
+            mul5,
+        )
+
+    sample_input = (
+        torch.rand((64,), dtype=torch.float32),
+        torch.rand((64,), dtype=torch.float32),
+    )
+
+    torch_model = BasicAddModelWithMultipleOutputs().eval()
+    edge_model = ai_edge_torch.convert(torch_model, sample_input)
+
+    result = model_coverage.compare_tflite_torch(edge_model, torch_model, sample_input)
+    self.assertTrue(result)
+
+  def test_apply_tfl_backdoor_flags(self):
+    """Tests if _apply_tfl_backdoor_flags correctly sets the values in a Converter object."""
+
+    class MockConverterInternalObject:
+
+      def __init__(self):
+        self.subkey2 = "original_subvalue2"
+
+    class MockConverter:
+
+      def __init__(self):
+        self.key1 = "original_value1"
+        self.key2 = MockConverterInternalObject()
+
+    mock_converter = MockConverter()
+    flags = {"key1": "new_value1", "key2": {"subkey2": "new_subvalue2"}}
+    cutils._apply_tfl_backdoor_flags(mock_converter, flags)
+
+    self.assertTrue(flags["key1"], "new_value1")
+    self.assertTrue(flags["key2"]["subkey2"], "new_subvalue2")
+
+  def test_convert_add_backdoor_flags(self):
+    """Tests conversion of an add module setting a tflite converter flag."""
+
+    class Add(torch.nn.Module):
+
+      def forward(self, a, b):
+        return a + b
+
+    args = (
+        torch.randn((5, 10)),
+        torch.randn((5, 10)),
+    )
+    torch_module = Add().eval()
+
+    with tempfile.TemporaryDirectory() as tmp_dir_path:
+      ir_dump_path = os.path.join(
+          tmp_dir_path, "test_convert_add_backdoor_flags_mlir_dump"
+      )
+      ai_edge_torch.convert(
+          torch_module, args, _ai_edge_converter_flags={"ir_dump_dir": ir_dump_path}
+      )
+      self.assertTrue(os.path.isdir(ir_dump_path))
+
+  def test_convert_model_with_dynamic_batch(self):
+    """
+    Test converting a simple model with dynamic batch size.
+    """
+
+    class SampleModel(torch.nn.Module):
+
+      def __init__(self):
+        super().__init__()
+        self.w = torch.ones((10, 10)) * 2.7
+
+      def forward(self, x, y):
+        return x + y + self.w
+
+    sample_input = (torch.randn(4, 3, 10, 10), torch.randn(4, 3, 10, 10))
+    batch = torch.export.Dim("batch")
+    dynamic_shapes = ({0: batch}, {0: batch})
+
+    model = SampleModel().eval()
+    edge_model = ai_edge_torch.convert(
+        model, sample_input, dynamic_shapes=dynamic_shapes
+    )
+
+    for batch_size in [2, 4, 10]:
+      validate_input = (
+          torch.randn(batch_size, 3, 10, 10),
+          torch.randn(batch_size, 3, 10, 10),
+      )
+      self.assertTrue(
+          model_coverage.compare_tflite_torch(edge_model, model, validate_input)
+      )
+
+  def test_convert_model_with_kwargs(self):
+    """
+    Test converting a simple model with sample_kwargs.
+    """
+
+    class SampleModel(torch.nn.Module):
+
+      def forward(self, x, y):
+        return x + y
+
+    kwargs_gen = lambda: dict(x=torch.randn(10, 10), y=torch.randn(10, 10))
+
+    model = SampleModel().eval()
+    edge_model = ai_edge_torch.convert(model, sample_kwargs=kwargs_gen())
+
+    self.assertTrue(
+        model_coverage.compare_tflite_torch(edge_model, model, kwargs=kwargs_gen)
+    )
+
+  def test_convert_model_with_args_kwargs(self):
+    """
+    Test converting a simple model with both sample_args and sample_kwargs.
+    """
+
+    class SampleModel(torch.nn.Module):
+
+      def forward(self, x, y):
+        return x + y
+
+    args_gen = lambda: (torch.randn(10, 10),)
+    kwargs_gen = lambda: dict(y=torch.randn(10, 10))
+
+    model = SampleModel().eval()
+    edge_model = ai_edge_torch.convert(model, args_gen(), kwargs_gen())
+
+    self.assertTrue(
+        model_coverage.compare_tflite_torch(edge_model, model, args_gen, kwargs_gen)
+    )
+
+
+if __name__ == "__main__":
+  unittest.main()