fix: remove USE_PYTHON_IMPL flag from nightly tests (#3678)

Signed-off-by: Kyunggeun Lee <[email protected]>

NightlyTests/torch/test_cross_layer_equalization.py

@@ -1,37 +1,37 @@
 # -*- mode: python -*-
 # =============================================================================
 #  @@-COPYRIGHT-START-@@
-#  
+#
 #  Copyright (c) 2017-2024, Qualcomm Innovation Center, Inc. All rights reserved.
-#  
-#  Redistribution and use in source and binary forms, with or without 
+#
+#  Redistribution and use in source and binary forms, with or without
 #  modification, are permitted provided that the following conditions are met:
-#  
-#  1. Redistributions of source code must retain the above copyright notice, 
+#
+#  1. Redistributions of source code must retain the above copyright notice,
 #     this list of conditions and the following disclaimer.
-#  
-#  2. Redistributions in binary form must reproduce the above copyright notice, 
-#     this list of conditions and the following disclaimer in the documentation 
+#
+#  2. Redistributions in binary form must reproduce the above copyright notice,
+#     this list of conditions and the following disclaimer in the documentation
 #     and/or other materials provided with the distribution.
-#  
-#  3. Neither the name of the copyright holder nor the names of its contributors 
-#     may be used to endorse or promote products derived from this software 
+#
+#  3. Neither the name of the copyright holder nor the names of its contributors
+#     may be used to endorse or promote products derived from this software
 #     without specific prior written permission.
-#  
-#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 
-#  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 
-#  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE 
-#  ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE 
-#  LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR 
-#  CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF 
-#  SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS 
-#  INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN 
-#  CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) 
+#
+#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+#  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+#  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+#  ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+#  LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+#  CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+#  SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+#  INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+#  CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 #  ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 #  POSSIBILITY OF SUCH DAMAGE.
-#  
+#
 #  SPDX-License-Identifier: BSD-3-Clause
-#  
+#
 #  @@-COPYRIGHT-END-@@
 # =============================================================================
 """ Cross Layer Equalization acceptance tests for ResNet model. """
@@ -42,38 +42,18 @@
 import copy
 import torch
 import numpy as np
-from contextlib import contextmanager
 from torchvision import models
 
-from aimet_torch import cross_layer_equalization as cle
 from aimet_torch import batch_norm_fold
 from aimet_torch.cross_layer_equalization import CrossLayerScaling, HighBiasFold, equalize_model
 from aimet_torch import visualize_model
 from models.mobilenet import MobileNetV2
 
 
-@contextmanager
-def _use_python_impl(flag: bool):
-    orig_flag = cle.USE_PYTHON_IMPL
-    try:
-        cle.USE_PYTHON_IMPL = flag
-        yield
-    finally:
-        cle.USE_PYTHON_IMPL = orig_flag
-
-
-@pytest.fixture(params=[True, False])
-def use_python_impl(request):
-    param: bool = request.param
-
-    with _use_python_impl(param):
-        yield
-
-
 class TestCrossLayerEqualization:
     """ Acceptance tests related to winnowing ResNet models. """
 
-    def test_cross_layer_equalization_resnet(self, use_python_impl):
+    def test_cross_layer_equalization_resnet(self):
 
         torch.manual_seed(10)
         model = models.resnet18().eval()
@@ -107,20 +87,20 @@ def test_cross_layer_equalization_resnet(self, use_python_impl):
         for i in range(len(model.layer1[1].conv2.bias)):
             assert model.layer1[1].conv2.bias[i] <= b2[i]
 
-    def test_cross_layer_equalization_mobilenet_v2(self, use_python_impl):
+    def test_cross_layer_equalization_mobilenet_v2(self):
         torch.manual_seed(10)
         model = MobileNetV2().to(torch.device('cpu'))
         model = model.eval()
         equalize_model(model, (1, 3, 224, 224))
 
-    def test_cross_layer_equalization_vgg(self, use_python_impl):
+    def test_cross_layer_equalization_vgg(self):
         torch.manual_seed(10)
         model = models.vgg16().to(torch.device('cpu'))
         model = model.eval()
         equalize_model(model, (1, 3, 224, 224))
 
     @pytest.mark.skip("Takes 1 min 42 secs to run")
-    def test_cross_layer_equalization_mobilenet_v2_visualize_after_optimization(self, use_python_impl):
+    def test_cross_layer_equalization_mobilenet_v2_visualize_after_optimization(self):
         torch.manual_seed(10)
         model = MobileNetV2().to(torch.device('cpu'))
         model = model.eval()
@@ -134,7 +114,7 @@ def test_cross_layer_equalization_mobilenet_v2_visualize_after_optimization(self
         equalize_model(model, (1, 3, 224, 224))
         visualize_model.visualize_changes_after_optimization(model_copy, model, results_dir)
 
-    def test_cross_layer_equalization_resnet18_visualize_to_identify_problem_layers(self, use_python_impl):
+    def test_cross_layer_equalization_resnet18_visualize_to_identify_problem_layers(self):
         torch.manual_seed(10)
         model = models.resnet18().eval()
 
@@ -146,7 +126,7 @@ def test_cross_layer_equalization_resnet18_visualize_to_identify_problem_layers(
             visualize_model.visualize_relative_weight_ranges_to_identify_problematic_layers(model, tmp_dir)
             assert os.path.isfile(file)
 
-    def test_cle_transposed_conv2D(self, use_python_impl):
+    def test_cle_transposed_conv2D(self):
         class TransposedConvModel(torch.nn.Module):
             def __init__(self):
                 super(TransposedConvModel, self).__init__()
@@ -190,7 +170,7 @@ def forward(self, x):
         output_after_cle = model(random_input).detach().numpy()
         assert np.allclose(output_before_cle, output_after_cle, rtol=1.e-2)
 
-    def test_cle_depthwise_transposed_conv2D(self, use_python_impl):
+    def test_cle_depthwise_transposed_conv2D(self):
 
         class TransposedConvModel(torch.nn.Module):
             def __init__(self):
@@ -241,7 +221,7 @@ def forward(self, x):
         output_after_cle = model(random_input).detach().numpy()
         assert np.allclose(output_before_cle, output_after_cle, rtol=1.e-2)
 
-    def test_cle_for_maskrcnn(self, use_python_impl):
+    def test_cle_for_maskrcnn(self):
         class JITTraceableWrapper(torch.nn.Module):
             def __init__(self, model):
                 super().__init__()