[UT] move compression ut to an additional folder & speed up simulated annealing search (#4357)

nni/algorithms/compression/v2/pytorch/pruning/tools/task_generator.py

@@ -229,8 +229,11 @@ class SimulatedAnnealingTaskGenerator(TaskGenerator):
         if target_sparsity == 0:
             return [], []
 
+        low_limit = 0
         while True:
-            random_sparsity = sorted(np.random.uniform(0, 1, len(op_names)))
+            # This is to speed up finding the legal sparsity.
+            low_limit = (1 - low_limit) * 0.05 + low_limit
+            random_sparsity = sorted(np.random.uniform(low_limit, 1, len(op_names)))
             rescaled_sparsity = self._rescale_sparsity(random_sparsity, target_sparsity, op_names)
             if rescaled_sparsity is not None and rescaled_sparsity[0] >= 0 and rescaled_sparsity[-1] < 1:
                 break

pipelines/fast-test.yml

@@ -179,14 +179,14 @@ stages:
         set -e
         cd test
         python -m pytest ut --cov-config=.coveragerc \
-          --ignore=ut/sdk/test_pruners.py \
-          --ignore=ut/sdk/test_compressor_tf.py \
-          --ignore=ut/sdk/test_compressor_torch.py \
-          --ignore=ut/sdk/test_model_speedup.py
-        python -m pytest ut/sdk/test_pruners.py --cov-config=.coveragerc --cov-append
-        python -m pytest ut/sdk/test_compressor_tf.py --cov-config=.coveragerc --cov-append
-        python -m pytest ut/sdk/test_compressor_torch.py --cov-config=.coveragerc --cov-append
-        python -m pytest ut/sdk/test_model_speedup.py --cov-config=.coveragerc --cov-append
+          --ignore=ut/compression/v1/test_pruners.py \
+          --ignore=ut/compression/v1/test_compressor_tf.py \
+          --ignore=ut/compression/v1/test_compressor_torch.py \
+          --ignore=ut/compression/v1/test_model_speedup.py
+        python -m pytest ut/compression/v1/test_pruners.py --cov-config=.coveragerc --cov-append
+        python -m pytest ut/compression/v1/test_compressor_tf.py --cov-config=.coveragerc --cov-append
+        python -m pytest ut/compression/v1/test_compressor_torch.py --cov-config=.coveragerc --cov-append
+        python -m pytest ut/compression/v1/test_model_speedup.py --cov-config=.coveragerc --cov-append
         cp coverage.xml ../coverage/python.xml
       displayName: Python unit test
 

test/ut/sdk/test_pruners.py → test/ut/compression/v1/test_pruners.py

@@ -15,8 +15,8 @@ from nni.algorithms.compression.pytorch.pruning import LevelPruner, SlimPruner,
     TaylorFOWeightFilterPruner, NetAdaptPruner, SimulatedAnnealingPruner, ADMMPruner, \
     AutoCompressPruner, AMCPruner
 
-sys.path.append(os.path.dirname(__file__))
-from models.pytorch_models.mobilenet import MobileNet
+sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
+from sdk.models.pytorch_models.mobilenet import MobileNet
 
 def validate_sparsity(wrapper, sparsity, bias=False):
     masks = [wrapper.weight_mask]

test/ut/sdk/test_transformer_pruners.py → test/ut/compression/v1/test_transformer_pruners.py

@@ -13,8 +13,8 @@ from unittest import TestCase, main
 
 from nni.algorithms.compression.pytorch.pruning import TransformerHeadPruner
 
-sys.path.append(os.path.dirname(__file__))
-from models.pytorch_models.transformer import TransformerEncoder
+sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
+from sdk.models.pytorch_models.transformer import TransformerEncoder
 
 
 def validate_sparsity(wrapper, sparsity, bias=False):

test/ut/sdk/test_v2_iterative_pruner_torch.py → test/ut/compression/v2/test_iterative_pruner_torch.py

@@ -41,6 +41,7 @@ class TorchModel(torch.nn.Module):
 
 def trainer(model, optimizer, criterion):
     model.train()
+    for _ in range(10):
         input = torch.rand(10, 1, 28, 28)
         label = torch.Tensor(list(range(10))).type(torch.LongTensor)
         optimizer.zero_grad()
@@ -65,7 +66,7 @@ class IterativePrunerTestCase(unittest.TestCase):
     def test_linear_pruner(self):
         model = TorchModel()
         config_list = [{'op_types': ['Conv2d'], 'sparsity': 0.8}]
-        pruner = LinearPruner(model, config_list, 'level', 3, log_dir='../../logs')
+        pruner = LinearPruner(model, config_list, 'level', 3, log_dir='../../../logs')
         pruner.compress()
         _, pruned_model, masks, _, _ = pruner.get_best_result()
         sparsity_list = compute_sparsity_mask2compact(pruned_model, masks, config_list)
@@ -74,7 +75,7 @@ class IterativePrunerTestCase(unittest.TestCase):
     def test_agp_pruner(self):
         model = TorchModel()
         config_list = [{'op_types': ['Conv2d'], 'sparsity': 0.8}]
-        pruner = AGPPruner(model, config_list, 'level', 3, log_dir='../../logs')
+        pruner = AGPPruner(model, config_list, 'level', 3, log_dir='../../../logs')
         pruner.compress()
         _, pruned_model, masks, _, _ = pruner.get_best_result()
         sparsity_list = compute_sparsity_mask2compact(pruned_model, masks, config_list)
@@ -83,7 +84,7 @@ class IterativePrunerTestCase(unittest.TestCase):
     def test_lottery_ticket_pruner(self):
         model = TorchModel()
         config_list = [{'op_types': ['Conv2d'], 'sparsity': 0.8}]
-        pruner = LotteryTicketPruner(model, config_list, 'level', 3, log_dir='../../logs')
+        pruner = LotteryTicketPruner(model, config_list, 'level', 3, log_dir='../../../logs')
         pruner.compress()
         _, pruned_model, masks, _, _ = pruner.get_best_result()
         sparsity_list = compute_sparsity_mask2compact(pruned_model, masks, config_list)
@@ -92,7 +93,7 @@ class IterativePrunerTestCase(unittest.TestCase):
     def test_simulated_annealing_pruner(self):
         model = TorchModel()
         config_list = [{'op_types': ['Conv2d'], 'sparsity': 0.8}]
-        pruner = SimulatedAnnealingPruner(model, config_list, evaluator, start_temperature=40, log_dir='../../logs')
+        pruner = SimulatedAnnealingPruner(model, config_list, evaluator, start_temperature=40, log_dir='../../../logs')
         pruner.compress()
         _, pruned_model, masks, _, _ = pruner.get_best_result()
         sparsity_list = compute_sparsity_mask2compact(pruned_model, masks, config_list)
@@ -112,7 +113,7 @@ class IterativePrunerTestCase(unittest.TestCase):
             'evaluator': evaluator,
             'start_temperature': 40
         }
-        pruner = AutoCompressPruner(model, config_list, 10, admm_params, sa_params=sa_params, log_dir='../../logs')
+        pruner = AutoCompressPruner(model, config_list, 10, admm_params, sa_params=sa_params, log_dir='../../../logs')
         pruner.compress()
         _, pruned_model, masks, _, _ = pruner.get_best_result()
         sparsity_list = compute_sparsity_mask2compact(pruned_model, masks, config_list)

test/ut/sdk/test_v2_pruner_torch.py → test/ut/compression/v2/test_pruner_torch.py

@@ -44,6 +44,7 @@ class TorchModel(torch.nn.Module):
 
 def trainer(model, optimizer, criterion):
     model.train()
+    for _ in range(10):
         input = torch.rand(10, 1, 28, 28)
         label = torch.Tensor(list(range(10))).type(torch.LongTensor)
         optimizer.zero_grad()

test/ut/sdk/test_v2_pruning_tools_torch.py → test/ut/compression/v2/test_pruning_tools_torch.py

@@ -50,6 +50,7 @@ class TorchModel(torch.nn.Module):
 
 def trainer(model, optimizer, criterion):
     model.train()
+    for _ in range(10):
         input = torch.rand(10, 1, 28, 28)
         label = torch.Tensor(list(range(10))).type(torch.LongTensor)
         optimizer.zero_grad()

test/ut/sdk/test_v2_task_generator.py → test/ut/compression/v2/test_task_generator.py

@@ -7,7 +7,7 @@ import unittest
 import torch
 import torch.nn.functional as F
 
-from nni.algorithms.compression.v2.pytorch.base import Task, TaskResult
+from nni.algorithms.compression.v2.pytorch.base import TaskResult
 from nni.algorithms.compression.v2.pytorch.pruning.tools import (
     AGPTaskGenerator,
     LinearTaskGenerator,