From fb06ab18e416dc749f61d2f2dee23cc70c2f805f Mon Sep 17 00:00:00 2001
From: J-shang <shangning128@163.com>
Date: Fri, 20 Aug 2021 14:13:16 +0800
Subject: [PATCH 01/16] add pruning tools test

---
 .../pytorch/test/test_pruning_tools_torch.py  | 201 ++++++++++++++++++
 1 file changed, 201 insertions(+)
 create mode 100644 nni/algorithms/compression/v2/pytorch/test/test_pruning_tools_torch.py

diff --git a/nni/algorithms/compression/v2/pytorch/test/test_pruning_tools_torch.py b/nni/algorithms/compression/v2/pytorch/test/test_pruning_tools_torch.py
new file mode 100644
index 0000000000..5d61cba62c
--- /dev/null
+++ b/nni/algorithms/compression/v2/pytorch/test/test_pruning_tools_torch.py
@@ -0,0 +1,201 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+from nni.algorithms.compression.v2.pytorch.pruning.tools.base import HookCollectorInfo
+import unittest
+
+import torch
+import torch.nn.functional as F
+
+from nni.algorithms.compression.v2.pytorch.base import Pruner
+from nni.algorithms.compression.v2.pytorch.pruning.tools import (
+    WeightDataCollector,
+    WeightTrainerBasedDataCollector,
+    SingleHookTrainerBasedDataCollector
+)
+from nni.algorithms.compression.v2.pytorch.pruning.tools import (
+    NormMetricsCalculator,
+    MultiDataNormMetricsCalculator,
+    DistMetricsCalculator,
+    APoZRankMetricsCalculator,
+    MeanRankMetricsCalculator
+)
+from nni.algorithms.compression.v2.pytorch.pruning.tools import (
+    NormalSparsityAllocator,
+    GlobalSparsityAllocator
+)
+from nni.compression.pytorch.utils import get_module_by_name
+
+
+class TorchModel(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = torch.nn.Conv2d(1, 5, 5, 1)
+        self.bn1 = torch.nn.BatchNorm2d(5)
+        self.conv2 = torch.nn.Conv2d(5, 10, 5, 1)
+        self.bn2 = torch.nn.BatchNorm2d(10)
+        self.fc1 = torch.nn.Linear(4 * 4 * 10, 100)
+        self.fc2 = torch.nn.Linear(100, 10)
+
+    def forward(self, x):
+        x = F.relu(self.bn1(self.conv1(x)))
+        x = F.max_pool2d(x, 2, 2)
+        x = F.relu(self.bn2(self.conv2(x)))
+        x = F.max_pool2d(x, 2, 2)
+        x = x.view(-1, 4 * 4 * 10)
+        x = F.relu(self.fc1(x))
+        x = self.fc2(x)
+        return F.log_softmax(x, dim=1)
+
+
+def trainer(model, optimizer, criterion):
+    model.train()
+    input = torch.rand(10, 1, 28, 28)
+    label = torch.Tensor(list(range(10))).type(torch.LongTensor)
+    optimizer.zero_grad()
+    output = model(input)
+    loss = criterion(output, label)
+    loss.backward()
+    optimizer.step()
+
+
+def get_optimizer(model):
+    return torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4)
+
+
+criterion = torch.nn.CrossEntropyLoss()
+
+
+class PruningToolsTestCase(unittest.TestCase):
+    def test_data_collector(self):
+        model = TorchModel()
+        w1 = torch.rand(5, 1, 5, 5)
+        w2 = torch.rand(10, 5, 5, 5)
+        model.conv1.weight.data = w1
+        model.conv2.weight.data = w2
+
+        config_list = [{'op_types': ['Conv2d']}]
+        pruner = Pruner(model, config_list)
+
+        # Test WeightDataCollector
+        data_collector = WeightDataCollector(pruner)
+        data = data_collector.collect()
+        assert all(torch.equal(get_module_by_name(model, module_name)[1].module.weight.data, data[module_name]) for module_name in ['conv1', 'conv2'])
+
+        # Test WeightTrainerBasedDataCollector
+        def opt_after():
+            model.conv1.module.weight.data = torch.ones(5, 1, 5, 5)
+            model.conv2.module.weight.data = torch.ones(10, 5, 5, 5)
+
+        data_collector = WeightTrainerBasedDataCollector(pruner, trainer, get_optimizer(model), criterion, 1, opt_after_tasks=[opt_after])
+        data = data_collector.collect()
+        assert all(torch.equal(get_module_by_name(model, module_name)[1].module.weight.data, data[module_name]) for module_name in ['conv1', 'conv2'])
+        assert all(t.numel() == (t == 1).type_as(t).sum().item() for t in data.values())
+
+        # Test SingleHookTrainerBasedDataCollector
+        def _collector(buffer, weight_tensor):
+            def collect_taylor(grad):
+                if len(buffer) < 2:
+                    buffer.append(grad.clone().detach())
+            return collect_taylor
+        hook_targets = {'conv1': model.conv1.module.weight, 'conv2': model.conv2.module.weight}
+        collector_info = HookCollectorInfo(hook_targets, 'tensor', _collector)
+
+        data_collector = SingleHookTrainerBasedDataCollector(pruner, trainer, get_optimizer(model), criterion, 2, collector_infos=[collector_info])
+        data = data_collector.collect()
+        assert all(len(t) == 2 for t in data.values())
+
+    def test_metrics_calculator(self):
+        # Test NormMetricsCalculator
+        metrics_calculator = NormMetricsCalculator(dim=0, p=2)
+        data = {
+            '1': torch.ones(3, 3, 3),
+            '2': torch.ones(4, 4) * 2
+        }
+        result = {
+            '1': torch.ones(3) * 3,
+            '2': torch.ones(4) * 4
+        }
+        metrics = metrics_calculator.calculate_metrics(data)
+        assert all(torch.equal(result[k], v) for k, v in metrics.items())
+
+        # Test DistMetricsCalculator
+        metrics_calculator = DistMetricsCalculator(dim=0, p=2)
+        data = {
+            '1': torch.tensor([[1, 2], [4, 6]], dtype=torch.float32),
+            '2': torch.tensor([[0, 0], [1, 1]], dtype=torch.float32)
+        }
+        result = {
+            '1': torch.tensor([5, 5], dtype=torch.float32),
+            '2': torch.sqrt(torch.tensor([2, 2], dtype=torch.float32))
+        }
+        metrics = metrics_calculator.calculate_metrics(data)
+        assert all(torch.equal(result[k], v) for k, v in metrics.items())
+
+        # Test MultiDataNormMetricsCalculator
+        metrics_calculator = MultiDataNormMetricsCalculator(dim=0, p=1)
+        data = {
+            '1': [torch.ones(3, 3, 3), torch.ones(3, 3, 3) * 2],
+            '2': [torch.ones(4, 4), torch.ones(4, 4) * 2]
+        }
+        result = {
+            '1': torch.ones(3) * 27,
+            '2': torch.ones(4) * 12
+        }
+        metrics = metrics_calculator.calculate_metrics(data)
+        assert all(torch.equal(result[k], v) for k, v in metrics.items())
+
+        # Test APoZRankMetricsCalculator
+        metrics_calculator = APoZRankMetricsCalculator(dim=1)
+        data = {
+            '1': [torch.tensor([[1, 0], [0, 1]], dtype=torch.float32), torch.tensor([[0, 1], [1, 0]], dtype=torch.float32)],
+            '2': [torch.tensor([[1, 0, 1], [0, 1, 0]], dtype=torch.float32), torch.tensor([[0, 0, 1], [0, 0, 0]], dtype=torch.float32)]
+        }
+        result = {
+            '1': torch.tensor([0.5, 0.5], dtype=torch.float32),
+            '2': torch.tensor([0.25, 0.25, 0.5], dtype=torch.float32)
+        }
+        metrics = metrics_calculator.calculate_metrics(data)
+        assert all(torch.equal(result[k], v) for k, v in metrics.items())
+
+        # Test MeanRankMetricsCalculator
+        metrics_calculator = MeanRankMetricsCalculator(dim=1)
+        data = {
+            '1': [torch.tensor([[1, 0], [0, 1]], dtype=torch.float32), torch.tensor([[0, 1], [1, 0]], dtype=torch.float32)],
+            '2': [torch.tensor([[1, 0, 1], [0, 1, 0]], dtype=torch.float32), torch.tensor([[0, 0, 1], [0, 0, 0]], dtype=torch.float32)]
+        }
+        result = {
+            '1': torch.tensor([0.5, 0.5], dtype=torch.float32),
+            '2': torch.tensor([0.25, 0.25, 0.5], dtype=torch.float32)
+        }
+        metrics = metrics_calculator.calculate_metrics(data)
+        assert all(torch.equal(result[k], v) for k, v in metrics.items())
+
+    def test_sparsity_allocator(self):
+        # Test NormalSparsityAllocator
+        model = TorchModel()
+        config_list = [{'op_types': ['Conv2d'], 'total_sparsity': 0.8}]
+        pruner = Pruner(model, config_list)
+        metrics = {
+            'conv1': torch.rand(5, 1, 5, 5),
+            'conv2': torch.rand(10, 5, 5, 5)
+        }
+        sparsity_allocator = NormalSparsityAllocator(pruner)
+        masks = sparsity_allocator.generate_sparsity(metrics)
+        assert all(v['weight_mask'].sum() / v['weight_mask'].numel() == 0.2 for k, v in masks.items())
+
+        # Test GlobalSparsityAllocator
+        model = TorchModel()
+        config_list = [{'op_types': ['Conv2d'], 'total_sparsity': 0.8}]
+        pruner = Pruner(model, config_list)
+        sparsity_allocator = GlobalSparsityAllocator(pruner)
+        masks = sparsity_allocator.generate_sparsity(metrics)
+        total_elements, total_masked_elements = 0, 0
+        for t in masks.values():
+            total_elements += t['weight_mask'].numel()
+            total_masked_elements += t['weight_mask'].sum().item()
+        assert total_masked_elements / total_elements == 0.2
+
+
+if __name__ == '__main__':
+    unittest.main()

From fe2c09d36a0e836ff262e2b1cb4f9552ff4e685a Mon Sep 17 00:00:00 2001
From: J-shang <shangning128@163.com>
Date: Mon, 30 Aug 2021 14:22:05 +0800
Subject: [PATCH 02/16] add pruner test

---
 .../v2/pytorch/test/test_pruner_torch.py      | 125 ++++++++++++++++++
 1 file changed, 125 insertions(+)
 create mode 100644 nni/algorithms/compression/v2/pytorch/test/test_pruner_torch.py

diff --git a/nni/algorithms/compression/v2/pytorch/test/test_pruner_torch.py b/nni/algorithms/compression/v2/pytorch/test/test_pruner_torch.py
new file mode 100644
index 0000000000..5785229449
--- /dev/null
+++ b/nni/algorithms/compression/v2/pytorch/test/test_pruner_torch.py
@@ -0,0 +1,125 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+from logging import critical
+import unittest
+
+import torch
+import torch.nn.functional as F
+
+from nni.algorithms.compression.v2.pytorch.pruning import (
+    LevelPruner,
+    L1NormPruner,
+    L2NormPruner,
+    SlimPruner,
+    FPGMPruner,
+    ActivationAPoZRankPruner,
+    ActivationMeanRankPruner,
+    TaylorFOWeightPruner
+)
+
+
+class TorchModel(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = torch.nn.Conv2d(1, 5, 5, 1)
+        self.bn1 = torch.nn.BatchNorm2d(5)
+        self.conv2 = torch.nn.Conv2d(5, 10, 5, 1)
+        self.bn2 = torch.nn.BatchNorm2d(10)
+        self.fc1 = torch.nn.Linear(4 * 4 * 10, 100)
+        self.fc2 = torch.nn.Linear(100, 10)
+
+    def forward(self, x):
+        x = F.relu(self.bn1(self.conv1(x)))
+        x = F.max_pool2d(x, 2, 2)
+        x = F.relu(self.bn2(self.conv2(x)))
+        x = F.max_pool2d(x, 2, 2)
+        x = x.view(-1, 4 * 4 * 10)
+        x = F.relu(self.fc1(x))
+        x = self.fc2(x)
+        return F.log_softmax(x, dim=1)
+
+
+def trainer(model, optimizer, criterion):
+    model.train()
+    input = torch.rand(10, 1, 28, 28)
+    label = torch.Tensor(list(range(10))).type(torch.LongTensor)
+    optimizer.zero_grad()
+    output = model(input)
+    loss = criterion(output, label)
+    loss.backward()
+    optimizer.step()
+
+
+def get_optimizer(model):
+    return torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4)
+
+
+criterion = torch.nn.CrossEntropyLoss()
+
+
+class PrunerTestCase(unittest.TestCase):
+    def test_level_pruner(self):
+        model = TorchModel()
+        config_list = [{'op_types': ['Conv2d'], 'sparsity': 0.8}]
+        pruner = LevelPruner(model=model, config_list=config_list)
+        pruned_model, masks = pruner.compress()
+        # TODO: validate masks sparsity
+
+    def test_l1_norm_pruner(self):
+        model = TorchModel()
+        config_list = [{'op_types': ['Conv2d'], 'sparsity': 0.8}]
+        pruner = L1NormPruner(model=model, config_list=config_list, mode='dependency_aware',
+                              dummy_input=torch.rand(10, 1, 28, 28))
+        pruned_model, masks = pruner.compress()
+
+    def test_l2_norm_pruner(self):
+        model = TorchModel()
+        config_list = [{'op_types': ['Conv2d'], 'sparsity': 0.8}]
+        pruner = L2NormPruner(model=model, config_list=config_list, mode='dependency_aware',
+                              dummy_input=torch.rand(10, 1, 28, 28))
+        pruned_model, masks = pruner.compress()
+
+    def test_fpgm_pruner(self):
+        model = TorchModel()
+        config_list = [{'op_types': ['Conv2d'], 'sparsity': 0.8}]
+        pruner = FPGMPruner(model=model, config_list=config_list, mode='dependency_aware',
+                            dummy_input=torch.rand(10, 1, 28, 28))
+        pruned_model, masks = pruner.compress()
+
+    def test_slim_pruner(self):
+        model = TorchModel()
+        config_list = [{'op_types': ['BatchNorm2d'], 'total_sparsity': 0.8}]
+        pruner = SlimPruner(model=model, config_list=config_list, trainer=trainer, optimizer=get_optimizer(model),
+                            criterion=criterion, training_epochs=1, scale=0.001, mode='global')
+        pruned_model, masks = pruner.compress()
+
+    def test_activation_apoz_rank_pruner(self):
+        model = TorchModel()
+        config_list = [{'op_types': ['Conv2d'], 'sparsity': 0.8}]
+        pruner = ActivationAPoZRankPruner(model=model, config_list=config_list, trainer=trainer,
+                                          optimizer=get_optimizer(model), criterion=criterion, training_batches=1,
+                                          activation='relu', mode='dependency_aware',
+                                          dummy_input=torch.rand(10, 1, 28, 28))
+        pruned_model, masks = pruner.compress()
+
+    def test_activation_mean_rank_pruner(self):
+        model = TorchModel()
+        config_list = [{'op_types': ['Conv2d'], 'sparsity': 0.8}]
+        pruner = ActivationMeanRankPruner(model=model, config_list=config_list, trainer=trainer,
+                                          optimizer=get_optimizer(model), criterion=criterion, training_batches=1,
+                                          activation='relu', mode='dependency_aware',
+                                          dummy_input=torch.rand(10, 1, 28, 28))
+        pruned_model, masks = pruner.compress()
+
+    def test_taylor_fo_pruner(self):
+        model = TorchModel()
+        config_list = [{'op_types': ['Conv2d'], 'sparsity': 0.8}]
+        pruner = TaylorFOWeightPruner(model=model, config_list=config_list, trainer=trainer,
+                                      optimizer=get_optimizer(model), criterion=criterion, training_batches=1,
+                                      mode='dependency_aware', dummy_input=torch.rand(10, 1, 28, 28))
+        pruned_model, masks = pruner.compress()
+
+
+if __name__ == '__main__':
+    unittest.main()

From 97ffcf417b0455621a6ccadae48254d4c2712b83 Mon Sep 17 00:00:00 2001
From: J-shang <shangning128@163.com>
Date: Sat, 18 Sep 2021 11:10:47 +0800
Subject: [PATCH 03/16] fix bug

---
 .../compression/v2/pytorch/base/compressor.py |  2 +-
 .../pruning/tools/sparsity_allocator.py       |  4 +--
 .../pytorch/test/test_pruning_tools_torch.py  |  6 ++--
 .../compression/v2/pytorch/utils/pruning.py   | 29 +++++++++++++++++++
 .../pytorch/utils/shape_dependency.py         |  3 +-
 5 files changed, 37 insertions(+), 7 deletions(-)

diff --git a/nni/algorithms/compression/v2/pytorch/base/compressor.py b/nni/algorithms/compression/v2/pytorch/base/compressor.py
index 1062fcfb24..c2801d38f9 100644
--- a/nni/algorithms/compression/v2/pytorch/base/compressor.py
+++ b/nni/algorithms/compression/v2/pytorch/base/compressor.py
@@ -9,7 +9,7 @@
 from torch.nn import Module
 
 from nni.common.graph_utils import TorchModuleGraph
-from nni.compression.pytorch.utils import get_module_by_name
+from nni.algorithms.compression.v2.pytorch.utils.pruning import get_module_by_name
 
 _logger = logging.getLogger(__name__)
 
diff --git a/nni/algorithms/compression/v2/pytorch/pruning/tools/sparsity_allocator.py b/nni/algorithms/compression/v2/pytorch/pruning/tools/sparsity_allocator.py
index 12f3e68b62..f0f0e97961 100644
--- a/nni/algorithms/compression/v2/pytorch/pruning/tools/sparsity_allocator.py
+++ b/nni/algorithms/compression/v2/pytorch/pruning/tools/sparsity_allocator.py
@@ -103,8 +103,8 @@ def __init__(self, pruner: Pruner, dim: int, dummy_input: Any):
 
     def _get_dependency(self):
         graph = self.pruner.generate_graph(dummy_input=self.dummy_input)
-        self.channel_depen = ChannelDependency(traced_model=graph.trace).dependency_sets
-        self.group_depen = GroupDependency(traced_model=graph.trace).dependency_sets
+        self.channel_depen = ChannelDependency(model=self.pruner.bound_model, dummy_input=self.dummy_input, traced_model=graph.trace).dependency_sets
+        self.group_depen = GroupDependency(model=self.pruner.bound_model, dummy_input=self.dummy_input, traced_model=graph.trace).dependency_sets
 
     def generate_sparsity(self, metrics: Dict) -> Dict[str, Dict[str, Tensor]]:
         self._get_dependency()
diff --git a/nni/algorithms/compression/v2/pytorch/test/test_pruning_tools_torch.py b/nni/algorithms/compression/v2/pytorch/test/test_pruning_tools_torch.py
index 5d61cba62c..e68cca09b9 100644
--- a/nni/algorithms/compression/v2/pytorch/test/test_pruning_tools_torch.py
+++ b/nni/algorithms/compression/v2/pytorch/test/test_pruning_tools_torch.py
@@ -182,7 +182,7 @@ def test_sparsity_allocator(self):
         }
         sparsity_allocator = NormalSparsityAllocator(pruner)
         masks = sparsity_allocator.generate_sparsity(metrics)
-        assert all(v['weight_mask'].sum() / v['weight_mask'].numel() == 0.2 for k, v in masks.items())
+        assert all(v['weight'].sum() / v['weight'].numel() == 0.2 for k, v in masks.items())
 
         # Test GlobalSparsityAllocator
         model = TorchModel()
@@ -192,8 +192,8 @@ def test_sparsity_allocator(self):
         masks = sparsity_allocator.generate_sparsity(metrics)
         total_elements, total_masked_elements = 0, 0
         for t in masks.values():
-            total_elements += t['weight_mask'].numel()
-            total_masked_elements += t['weight_mask'].sum().item()
+            total_elements += t['weight'].numel()
+            total_masked_elements += t['weight'].sum().item()
         assert total_masked_elements / total_elements == 0.2
 
 
diff --git a/nni/algorithms/compression/v2/pytorch/utils/pruning.py b/nni/algorithms/compression/v2/pytorch/utils/pruning.py
index 63022cee59..1f453d8cde 100644
--- a/nni/algorithms/compression/v2/pytorch/utils/pruning.py
+++ b/nni/algorithms/compression/v2/pytorch/utils/pruning.py
@@ -194,3 +194,32 @@ def get_model_weights_numel(model: Module, config_list: List[Dict], masks: Dict[
             else:
                 model_weights_numel[module_name] = module.weight.data.numel()
     return model_weights_numel, masked_rate
+
+# FIXME: to avoid circular import, copy this function in this place
+def get_module_by_name(model, module_name):
+    """
+    Get a module specified by its module name
+
+    Parameters
+    ----------
+    model : pytorch model
+        the pytorch model from which to get its module
+    module_name : str
+        the name of the required module
+
+    Returns
+    -------
+    module, module
+        the parent module of the required module, the required module
+    """
+    name_list = module_name.split(".")
+    for name in name_list[:-1]:
+        if hasattr(model, name):
+            model = getattr(model, name)
+        else:
+            return None, None
+    if hasattr(model, name_list[-1]):
+        leaf_module = getattr(model, name_list[-1])
+        return model, leaf_module
+    else:
+        return None, None
diff --git a/nni/compression/pytorch/utils/shape_dependency.py b/nni/compression/pytorch/utils/shape_dependency.py
index 82aa82c50a..f972212a5a 100644
--- a/nni/compression/pytorch/utils/shape_dependency.py
+++ b/nni/compression/pytorch/utils/shape_dependency.py
@@ -6,6 +6,7 @@
 import torch
 import numpy as np
 from nni.compression.pytorch.compressor import PrunerModuleWrapper
+from nni.algorithms.compression.v2.pytorch.base import PrunerModuleWrapper as PrunerModuleWrapper_v2
 from .utils import get_module_by_name
 
 
@@ -390,7 +391,7 @@ def _get_conv_groups(self, node_group):
         """
         node_name = node_group.name
         _, leaf_module = get_module_by_name(self.model, node_name)
-        if isinstance(leaf_module, PrunerModuleWrapper):
+        if isinstance(leaf_module, (PrunerModuleWrapper, PrunerModuleWrapper_v2)):
             leaf_module = leaf_module.module
         assert isinstance(
             leaf_module, (torch.nn.Conv2d, torch.nn.ConvTranspose2d))

From 91f42086cd38ddad059557f5468f6655146e049f Mon Sep 17 00:00:00 2001
From: J-shang <shangning128@163.com>
Date: Wed, 22 Sep 2021 10:18:24 +0800
Subject: [PATCH 04/16] update example

---
 .../pruning/v2/naive_prune_torch.py           | 153 ------------------
 .../pruning/v2/scheduler_torch.py             |  79 +++++++++
 .../pruning/v2/simple_pruning_torch.py        |  83 ++++++++++
 3 files changed, 162 insertions(+), 153 deletions(-)
 delete mode 100644 examples/model_compress/pruning/v2/naive_prune_torch.py
 create mode 100644 examples/model_compress/pruning/v2/scheduler_torch.py
 create mode 100644 examples/model_compress/pruning/v2/simple_pruning_torch.py

diff --git a/examples/model_compress/pruning/v2/naive_prune_torch.py b/examples/model_compress/pruning/v2/naive_prune_torch.py
deleted file mode 100644
index 648e6f3119..0000000000
--- a/examples/model_compress/pruning/v2/naive_prune_torch.py
+++ /dev/null
@@ -1,153 +0,0 @@
-import argparse
-import logging
-from pathlib import Path
-
-import torch
-from torchvision import transforms, datasets
-
-from nni.algorithms.compression.v2.pytorch import pruning
-from nni.compression.pytorch import ModelSpeedup
-from examples.model_compress.models.cifar10.vgg import VGG
-
-logging.getLogger().setLevel(logging.DEBUG)
-
-
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-model = VGG().to(device)
-
-normalize = transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
-
-train_loader = torch.utils.data.DataLoader(
-    datasets.CIFAR10('./data', train=True, transform=transforms.Compose([
-        transforms.RandomHorizontalFlip(),
-        transforms.RandomCrop(32, 4),
-        transforms.ToTensor(),
-        normalize,
-    ]), download=True),
-    batch_size=128, shuffle=True)
-
-test_loader = torch.utils.data.DataLoader(
-    datasets.CIFAR10('./data', train=False, transform=transforms.Compose([
-        transforms.ToTensor(),
-        normalize,
-    ])),
-    batch_size=200, shuffle=False)
-criterion = torch.nn.CrossEntropyLoss()
-
-def trainer(model, optimizer, criterion, epoch=None):
-    model.train()
-    for batch_idx, (data, target) in enumerate(train_loader):
-        data, target = data.to(device), target.to(device)
-        optimizer.zero_grad()
-        output = model(data)
-        loss = criterion(output, target)
-        loss.backward()
-        optimizer.step()
-        if batch_idx % 100 == 0:
-            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
-                epoch, batch_idx * len(data), len(train_loader.dataset),
-                100. * batch_idx / len(train_loader), loss.item()))
-
-def evaluator(model):
-    model.eval()
-    criterion = torch.nn.NLLLoss()
-    test_loss = 0
-    correct = 0
-    with torch.no_grad():
-        for data, target in test_loader:
-            data, target = data.to(device), target.to(device)
-            output = model(data)
-            test_loss += criterion(output, target).item()
-            pred = output.argmax(dim=1, keepdim=True)
-            correct += pred.eq(target.view_as(pred)).sum().item()
-    test_loss /= len(test_loader.dataset)
-    acc = 100 * correct / len(test_loader.dataset)
-
-    print('Test Loss: {}  Accuracy: {}%\n'.format(
-        test_loss, acc))
-    return acc
-
-optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4)
-fintune_optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4)
-
-def main(args):
-    if args.pre_train:
-        for i in range(1):
-            trainer(model, fintune_optimizer, criterion, epoch=i)
-
-    config_list = [{
-        'op_types': ['Conv2d'],
-        'sparsity_per_layer': 0.8
-    }]
-    kwargs = {
-        'model': model,
-        'config_list': config_list,
-    }
-    if args.pruner == 'level':
-        pruner = pruning.LevelPruner(**kwargs)
-    else:
-        kwargs['mode'] = args.mode
-        if kwargs['mode'] == 'dependency_aware':
-            kwargs['dummy_input'] = torch.rand(10, 3, 32, 32).to(device)
-        if args.pruner == 'l1norm':
-            pruner = pruning.L1NormPruner(**kwargs)
-        elif args.pruner == 'l2norm':
-            pruner = pruning.L2NormPruner(**kwargs)
-        elif args.pruner == 'fpgm':
-            pruner = pruning.FPGMPruner(**kwargs)
-        else:
-            kwargs['trainer'] = trainer
-            kwargs['optimizer'] = optimizer
-            kwargs['criterion'] = criterion
-            if args.pruner == 'slim':
-                kwargs['config_list'] = [{
-                    'op_types': ['BatchNorm2d'],
-                    'total_sparsity': 0.8,
-                    'max_sparsity_per_layer': 0.9
-                }]
-                kwargs['training_epochs'] = 1
-                pruner = pruning.SlimPruner(**kwargs)
-            elif args.pruner == 'mean_activation':
-                pruner = pruning.ActivationMeanRankPruner(**kwargs)
-            elif args.pruner == 'apoz':
-                pruner = pruning.ActivationAPoZRankPruner(**kwargs)
-            elif args.pruner == 'taylorfo':
-                pruner = pruning.TaylorFOWeightPruner(**kwargs)
-
-    pruned_model, masks = pruner.compress()
-    pruner.show_pruned_weights()
-
-    if args.speed_up:
-        tmp_masks = {}
-        for name, mask in masks.items():
-            tmp_masks[name] = {}
-            tmp_masks[name]['weight'] = mask.get('weight_mask')
-            if 'bias' in masks:
-                tmp_masks[name]['bias'] = mask.get('bias_mask')
-        torch.save(tmp_masks, Path('./temp_masks.pth'))
-        pruner._unwrap_model()
-        ModelSpeedup(model, torch.rand(10, 3, 32, 32).to(device), Path('./temp_masks.pth'))
-
-    if args.finetune:
-        for i in range(1):
-            trainer(pruned_model, fintune_optimizer, criterion, epoch=i)
-
-
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser(description='PyTorch CIFAR10 Example for model comporession')
-    parser.add_argument('--pruner', type=str, default='l1norm',
-                        choices=['level', 'l1norm', 'l2norm', 'slim',
-                                 'fpgm', 'mean_activation', 'apoz', 'taylorfo'],
-                        help='pruner to use')
-    parser.add_argument('--mode', type=str, default='normal',
-                        choices=['normal', 'dependency_aware', 'global'])
-    parser.add_argument('--pre-train', action='store_true', default=False,
-                        help='Whether to pre-train the model')
-    parser.add_argument('--speed-up', action='store_true', default=False,
-                        help='Whether to speed-up the pruned model')
-    parser.add_argument('--finetune', action='store_true', default=False,
-                        help='Whether to finetune the pruned model')
-    args = parser.parse_args()
-
-    main(args)
diff --git a/examples/model_compress/pruning/v2/scheduler_torch.py b/examples/model_compress/pruning/v2/scheduler_torch.py
new file mode 100644
index 0000000000..cb31db1f11
--- /dev/null
+++ b/examples/model_compress/pruning/v2/scheduler_torch.py
@@ -0,0 +1,79 @@
+import functools
+from tqdm import tqdm
+
+import torch
+from torchvision import datasets, transforms
+
+from nni.algorithms.compression.v2.pytorch.pruning import L1NormPruner
+from nni.algorithms.compression.v2.pytorch.pruning.tools import AGPTaskGenerator
+from nni.algorithms.compression.v2.pytorch.pruning.basic_scheduler import PruningScheduler
+
+from examples.model_compress.models.cifar10.vgg import VGG
+
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+normalize = transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
+
+train_loader = torch.utils.data.DataLoader(
+    datasets.CIFAR10('./data', train=True, transform=transforms.Compose([
+        transforms.RandomHorizontalFlip(),
+        transforms.RandomCrop(32, 4),
+        transforms.ToTensor(),
+        normalize,
+    ]), download=True),
+    batch_size=128, shuffle=True)
+
+test_loader = torch.utils.data.DataLoader(
+    datasets.CIFAR10('./data', train=False, transform=transforms.Compose([
+        transforms.ToTensor(),
+        normalize,
+    ])),
+    batch_size=128, shuffle=False)
+criterion = torch.nn.CrossEntropyLoss()
+
+def trainer(model, optimizer, criterion, epoch):
+    model.train()
+    for data, target in tqdm(iterable=train_loader, desc='Epoch {}'.format(epoch)):
+        data, target = data.to(device), target.to(device)
+        optimizer.zero_grad()
+        output = model(data)
+        loss = criterion(output, target)
+        loss.backward()
+        optimizer.step()
+
+def evaluator(model):
+    model.eval()
+    correct = 0
+    with torch.no_grad():
+        for data, target in tqdm(iterable=test_loader, desc='Test'):
+            data, target = data.to(device), target.to(device)
+            output = model(data)
+            pred = output.argmax(dim=1, keepdim=True)
+            correct += pred.eq(target.view_as(pred)).sum().item()
+    acc = 100 * correct / len(test_loader.dataset)
+    print('Accuracy: {}%\n'.format(acc))
+    return acc
+
+
+if __name__ == '__main__':
+    model = VGG().to(device)
+    optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4)
+    criterion = torch.nn.CrossEntropyLoss()
+
+    for i in range(5):
+        trainer(model, optimizer, criterion, i)
+
+    finetuner = functools.partial(trainer,
+                                  optimizer=torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4),
+                                  criterion=criterion,
+                                  epoch='PF')
+
+    config_list = [{'op_types': ['Conv2d'], 'sparsity': 0.8}]
+    task_generator = AGPTaskGenerator(10, model, config_list, keep_intermidiate_result=True)
+    pruner = L1NormPruner(model, config_list)
+
+    dummy_input = torch.rand(10, 3, 32, 32).to(device)
+    scheduler = PruningScheduler(pruner, task_generator, finetuner=finetuner, speed_up=True, dummy_input=dummy_input, evaluator=evaluator)
+
+    scheduler.compress()
diff --git a/examples/model_compress/pruning/v2/simple_pruning_torch.py b/examples/model_compress/pruning/v2/simple_pruning_torch.py
new file mode 100644
index 0000000000..f4c291d8fa
--- /dev/null
+++ b/examples/model_compress/pruning/v2/simple_pruning_torch.py
@@ -0,0 +1,83 @@
+from tqdm import tqdm
+
+import torch
+from torchvision import datasets, transforms
+
+from nni.algorithms.compression.v2.pytorch.pruning import L1NormPruner
+from nni.compression.pytorch.speedup import ModelSpeedup
+
+from examples.model_compress.models.cifar10.vgg import VGG
+
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+normalize = transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
+
+train_loader = torch.utils.data.DataLoader(
+    datasets.CIFAR10('./data', train=True, transform=transforms.Compose([
+        transforms.RandomHorizontalFlip(),
+        transforms.RandomCrop(32, 4),
+        transforms.ToTensor(),
+        normalize,
+    ]), download=True),
+    batch_size=128, shuffle=True)
+
+test_loader = torch.utils.data.DataLoader(
+    datasets.CIFAR10('./data', train=False, transform=transforms.Compose([
+        transforms.ToTensor(),
+        normalize,
+    ])),
+    batch_size=128, shuffle=False)
+criterion = torch.nn.CrossEntropyLoss()
+
+def trainer(model, optimizer, criterion, epoch):
+    model.train()
+    for data, target in tqdm(iterable=train_loader, desc='Epoch {}'.format(epoch)):
+        data, target = data.to(device), target.to(device)
+        optimizer.zero_grad()
+        output = model(data)
+        loss = criterion(output, target)
+        loss.backward()
+        optimizer.step()
+
+def evaluator(model):
+    model.eval()
+    correct = 0
+    with torch.no_grad():
+        for data, target in tqdm(iterable=test_loader, desc='Test'):
+            data, target = data.to(device), target.to(device)
+            output = model(data)
+            pred = output.argmax(dim=1, keepdim=True)
+            correct += pred.eq(target.view_as(pred)).sum().item()
+    acc = 100 * correct / len(test_loader.dataset)
+    print('Accuracy: {}%\n'.format(acc))
+    return acc
+
+
+if __name__ == '__main__':
+    model = VGG().to(device)
+    optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4)
+    criterion = torch.nn.CrossEntropyLoss()
+
+    print('\nPre-train the model:')
+    for i in range(5):
+        trainer(model, optimizer, criterion, i)
+        evaluator(model)
+
+    config_list = [{'op_types': ['Conv2d'], 'sparsity': 0.8}]
+    pruner = L1NormPruner(model, config_list)
+    _, masks = pruner.compress()
+
+    print('\nThe accuracy with masks:')
+    evaluator(model)
+
+    pruner._unwrap_model()
+    ModelSpeedup(model, dummy_input=torch.rand(10, 3, 32, 32).to(device), masks_file='simple_masks.pth').speedup_model()
+
+    print('\nThe accuracy after speed up:')
+    evaluator(model)
+
+    print('\nFinetune the model after speed up:')
+    for i in range(5):
+        trainer(model, optimizer, criterion, i)
+        evaluator(model)

From 1c13540be1bdf838d27cb260f22e3eef417d2e40 Mon Sep 17 00:00:00 2001
From: J-shang <shangning128@163.com>
Date: Wed, 22 Sep 2021 13:11:25 +0800
Subject: [PATCH 05/16] update unit test

---
 .../v2/pytorch/pruning/__init__.py            |  2 +
 .../v2/pytorch/test/test_pruner_torch.py      |  1 -
 .../pytorch/test/test_pruning_tools_torch.py  |  4 +-
 .../v2/pytorch/test/test_scheduler.py         | 46 ++++++++++++++
 .../v2/pytorch/test/test_task_generator.py    | 60 +++++++++++++++++++
 5 files changed, 110 insertions(+), 3 deletions(-)
 create mode 100644 nni/algorithms/compression/v2/pytorch/test/test_scheduler.py
 create mode 100644 nni/algorithms/compression/v2/pytorch/test/test_task_generator.py

diff --git a/nni/algorithms/compression/v2/pytorch/pruning/__init__.py b/nni/algorithms/compression/v2/pytorch/pruning/__init__.py
index d745ecf8d1..0c75ee896f 100644
--- a/nni/algorithms/compression/v2/pytorch/pruning/__init__.py
+++ b/nni/algorithms/compression/v2/pytorch/pruning/__init__.py
@@ -1 +1,3 @@
 from .basic_pruner import *
+from .basic_scheduler import PruningScheduler
+from .tools import AGPTaskGenerator, LinearTaskGenerator
diff --git a/nni/algorithms/compression/v2/pytorch/test/test_pruner_torch.py b/nni/algorithms/compression/v2/pytorch/test/test_pruner_torch.py
index 5785229449..780b1f503b 100644
--- a/nni/algorithms/compression/v2/pytorch/test/test_pruner_torch.py
+++ b/nni/algorithms/compression/v2/pytorch/test/test_pruner_torch.py
@@ -1,7 +1,6 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 
-from logging import critical
 import unittest
 
 import torch
diff --git a/nni/algorithms/compression/v2/pytorch/test/test_pruning_tools_torch.py b/nni/algorithms/compression/v2/pytorch/test/test_pruning_tools_torch.py
index e68cca09b9..579dd6f678 100644
--- a/nni/algorithms/compression/v2/pytorch/test/test_pruning_tools_torch.py
+++ b/nni/algorithms/compression/v2/pytorch/test/test_pruning_tools_torch.py
@@ -1,7 +1,6 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 
-from nni.algorithms.compression.v2.pytorch.pruning.tools.base import HookCollectorInfo
 import unittest
 
 import torch
@@ -24,7 +23,8 @@
     NormalSparsityAllocator,
     GlobalSparsityAllocator
 )
-from nni.compression.pytorch.utils import get_module_by_name
+from nni.algorithms.compression.v2.pytorch.pruning.tools.base import HookCollectorInfo
+from nni.algorithms.compression.v2.pytorch.utils.pruning import get_module_by_name
 
 
 class TorchModel(torch.nn.Module):
diff --git a/nni/algorithms/compression/v2/pytorch/test/test_scheduler.py b/nni/algorithms/compression/v2/pytorch/test/test_scheduler.py
new file mode 100644
index 0000000000..9b9b97a364
--- /dev/null
+++ b/nni/algorithms/compression/v2/pytorch/test/test_scheduler.py
@@ -0,0 +1,46 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+import unittest
+
+import torch
+import torch.nn.functional as F
+
+from nni.algorithms.compression.v2.pytorch.pruning import PruningScheduler, L1NormPruner, AGPTaskGenerator
+
+
+class TorchModel(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = torch.nn.Conv2d(1, 5, 5, 1)
+        self.bn1 = torch.nn.BatchNorm2d(5)
+        self.conv2 = torch.nn.Conv2d(5, 10, 5, 1)
+        self.bn2 = torch.nn.BatchNorm2d(10)
+        self.fc1 = torch.nn.Linear(4 * 4 * 10, 100)
+        self.fc2 = torch.nn.Linear(100, 10)
+
+    def forward(self, x):
+        x = F.relu(self.bn1(self.conv1(x)))
+        x = F.max_pool2d(x, 2, 2)
+        x = F.relu(self.bn2(self.conv2(x)))
+        x = F.max_pool2d(x, 2, 2)
+        x = x.view(-1, 4 * 4 * 10)
+        x = F.relu(self.fc1(x))
+        x = self.fc2(x)
+        return F.log_softmax(x, dim=1)
+
+
+class PruningSchedulerTestCase(unittest.TestCase):
+    def test_pruning_scheduler(self):
+        model = TorchModel()
+        config_list = [{'op_types': ['Conv2d'], 'sparsity': 0.8}]
+
+        task_generator = AGPTaskGenerator(1, model, config_list)
+        pruner = L1NormPruner(model, config_list)
+        scheduler = PruningScheduler(pruner, task_generator)
+
+        scheduler.compress()
+
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/nni/algorithms/compression/v2/pytorch/test/test_task_generator.py b/nni/algorithms/compression/v2/pytorch/test/test_task_generator.py
new file mode 100644
index 0000000000..efe7cc04d3
--- /dev/null
+++ b/nni/algorithms/compression/v2/pytorch/test/test_task_generator.py
@@ -0,0 +1,60 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+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.pruning import (
+    AGPTaskGenerator,
+    LinearTaskGenerator
+)
+
+
+class TorchModel(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = torch.nn.Conv2d(1, 5, 5, 1)
+        self.bn1 = torch.nn.BatchNorm2d(5)
+        self.conv2 = torch.nn.Conv2d(5, 10, 5, 1)
+        self.bn2 = torch.nn.BatchNorm2d(10)
+        self.fc1 = torch.nn.Linear(4 * 4 * 10, 100)
+        self.fc2 = torch.nn.Linear(100, 10)
+
+    def forward(self, x):
+        x = F.relu(self.bn1(self.conv1(x)))
+        x = F.max_pool2d(x, 2, 2)
+        x = F.relu(self.bn2(self.conv2(x)))
+        x = F.max_pool2d(x, 2, 2)
+        x = x.view(-1, 4 * 4 * 10)
+        x = F.relu(self.fc1(x))
+        x = self.fc2(x)
+        return F.log_softmax(x, dim=1)
+
+
+class TaskGenerator(unittest.TestCase):
+    def test_agp_task_generator(self):
+        model = TorchModel()
+        config_list = [{'op_types': ['Conv2d'], 'sparsity': 0.8}]
+
+        task_generator = AGPTaskGenerator(1, model, config_list)
+        tasks = task_generator.generate_tasks()
+
+        # TODO: wait for pr 4178
+        pass
+
+    def test_linear_task_generator(self):
+        model = TorchModel()
+        config_list = [{'op_types': ['Conv2d'], 'sparsity': 0.8}]
+
+        task_generator = LinearTaskGenerator(1, model, config_list)
+        tasks = task_generator.generate_tasks()
+
+        # TODO: wait for pr 4178
+        pass
+
+
+if __name__ == '__main__':
+    unittest.main()

From bfb7dfb966f9d0fc8575d234a8941fd6f394a1ec Mon Sep 17 00:00:00 2001
From: J-shang <shangning128@163.com>
Date: Sun, 26 Sep 2021 11:22:22 +0800
Subject: [PATCH 06/16] update test

---
 .../v2/pytorch/pruning/__init__.py            |  2 +-
 .../v2/pytorch/test/test_task_generator.py    | 63 ++++++++++++++-----
 2 files changed, 49 insertions(+), 16 deletions(-)

diff --git a/nni/algorithms/compression/v2/pytorch/pruning/__init__.py b/nni/algorithms/compression/v2/pytorch/pruning/__init__.py
index 0c75ee896f..19d1f11445 100644
--- a/nni/algorithms/compression/v2/pytorch/pruning/__init__.py
+++ b/nni/algorithms/compression/v2/pytorch/pruning/__init__.py
@@ -1,3 +1,3 @@
 from .basic_pruner import *
 from .basic_scheduler import PruningScheduler
-from .tools import AGPTaskGenerator, LinearTaskGenerator
+from .tools import AGPTaskGenerator, LinearTaskGenerator, LotteryTicketTaskGenerator, SimulatedAnnealingTaskGenerator
diff --git a/nni/algorithms/compression/v2/pytorch/test/test_task_generator.py b/nni/algorithms/compression/v2/pytorch/test/test_task_generator.py
index efe7cc04d3..8e98308104 100644
--- a/nni/algorithms/compression/v2/pytorch/test/test_task_generator.py
+++ b/nni/algorithms/compression/v2/pytorch/test/test_task_generator.py
@@ -1,6 +1,7 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 
+from typing import List
 import unittest
 
 import torch
@@ -9,7 +10,9 @@
 from nni.algorithms.compression.v2.pytorch.base import Task, TaskResult
 from nni.algorithms.compression.v2.pytorch.pruning import (
     AGPTaskGenerator,
-    LinearTaskGenerator
+    LinearTaskGenerator,
+    LotteryTicketTaskGenerator,
+    SimulatedAnnealingTaskGenerator
 )
 
 
@@ -34,26 +37,56 @@ def forward(self, x):
         return F.log_softmax(x, dim=1)
 
 
-class TaskGenerator(unittest.TestCase):
-    def test_agp_task_generator(self):
-        model = TorchModel()
-        config_list = [{'op_types': ['Conv2d'], 'sparsity': 0.8}]
+def test_task_generator(task_generator_type):
+    model = TorchModel()
+    config_list = [{'op_types': ['Conv2d'], 'sparsity': 0.8}]
+
+    if task_generator_type == 'agp':
+        task_generator = AGPTaskGenerator(5, model, config_list)
+    elif task_generator_type == 'linear':
+        task_generator = LinearTaskGenerator(5, model, config_list)
+    elif task_generator_type == 'lottery_ticket':
+        task_generator = LotteryTicketTaskGenerator(5, model, config_list)
+    elif task_generator_type == 'simulated_annealing':
+        task_generator = SimulatedAnnealingTaskGenerator(model, config_list)
+
+    count = run_task_generator(task_generator)
+
+    if task_generator_type == 'agp':
+        assert count == 6
+    elif task_generator_type == 'linear':
+        assert count == 6
+    elif task_generator_type == 'lottery_ticket':
+        assert count == 6
+    elif task_generator_type == 'simulated_annealing':
+        assert count == 17
 
-        task_generator = AGPTaskGenerator(1, model, config_list)
-        tasks = task_generator.generate_tasks()
 
-        # TODO: wait for pr 4178
-        pass
+def run_task_generator(task_generator):
+    task = task_generator.next()
+    factor = 0.9
+    count = 0
+    while task is not None:
+        factor = factor ** 2
+        count += 1
+        task_result = TaskResult(task.task_id, TorchModel(), {}, {}, 1 - factor)
+        task_generator.receive_task_result(task_result)
+        task = task_generator.next()
+    return count
+
+
+class TaskGenerator(unittest.TestCase):
+    def test_agp_task_generator(self):
+        test_task_generator('agp')
 
     def test_linear_task_generator(self):
-        model = TorchModel()
-        config_list = [{'op_types': ['Conv2d'], 'sparsity': 0.8}]
+        test_task_generator('linear')
 
-        task_generator = LinearTaskGenerator(1, model, config_list)
-        tasks = task_generator.generate_tasks()
+    def test_lottery_ticket_task_generator(self):
+        test_task_generator('lottery_ticket')
 
-        # TODO: wait for pr 4178
-        pass
+    def test_simulated_annealing_task_generator(self):
+        test_task_generator('simulated_annealing')
 
 
 if __name__ == '__main__':

From 3b2c655032ea9a39d8e8790b645b99afcf13e9f8 Mon Sep 17 00:00:00 2001
From: J-shang <shangning128@163.com>
Date: Sun, 26 Sep 2021 13:28:59 +0800
Subject: [PATCH 07/16] update pruner test

---
 .../v2/pytorch/pruning/basic_pruner.py        |  6 +--
 .../v2/pytorch/test/test_pruner_torch.py      | 38 ++++++++++++++++++-
 2 files changed, 39 insertions(+), 5 deletions(-)

diff --git a/nni/algorithms/compression/v2/pytorch/pruning/basic_pruner.py b/nni/algorithms/compression/v2/pytorch/pruning/basic_pruner.py
index a1d21ddd48..8157b28ce7 100644
--- a/nni/algorithms/compression/v2/pytorch/pruning/basic_pruner.py
+++ b/nni/algorithms/compression/v2/pytorch/pruning/basic_pruner.py
@@ -43,7 +43,7 @@
 _logger = logging.getLogger(__name__)
 
 __all__ = ['LevelPruner', 'L1NormPruner', 'L2NormPruner', 'FPGMPruner', 'SlimPruner', 'ActivationPruner',
-           'ActivationAPoZRankPruner', 'ActivationMeanRankPruner', 'TaylorFOWeightPruner']
+           'ActivationAPoZRankPruner', 'ActivationMeanRankPruner', 'TaylorFOWeightPruner', 'ADMMPruner']
 
 NORMAL_SCHEMA = {
     Or('sparsity', 'sparsity_per_layer'): And(float, lambda n: 0 <= n < 1),
@@ -677,7 +677,7 @@ def __init__(self, model: Module, config_list: List[Dict], trainer: Callable[[Mo
             Supported keys:
                 - sparsity : This is to specify the sparsity for each layer in this config to be compressed.
                 - sparsity_per_layer : Equals to sparsity.
-                - rho : Penalty parameters in ADMM algorithm.
+                - rho : Penalty parameters in ADMM algorithm. Default: 1e-4.
                 - op_types : Operation types to prune.
                 - op_names : Operation names to prune.
                 - exclude : Set True then the layers setting by op_types and op_names will be excluded from pruning.
@@ -732,7 +732,7 @@ def criterion_patch(self, origin_criterion: Callable[[Tensor, Tensor], Tensor]):
         def patched_criterion(output: Tensor, target: Tensor):
             penalty = torch.tensor(0.0).to(output.device)
             for name, wrapper in self.get_modules_wrapper().items():
-                rho = wrapper.config['rho']
+                rho = wrapper.config.get('rho', 1e-4)
                 penalty += (rho / 2) * torch.sqrt(torch.norm(wrapper.module.weight - self.Z[name] + self.U[name]))
             return origin_criterion(output, target) + penalty
         return patched_criterion
diff --git a/nni/algorithms/compression/v2/pytorch/test/test_pruner_torch.py b/nni/algorithms/compression/v2/pytorch/test/test_pruner_torch.py
index 780b1f503b..a727ecaeb8 100644
--- a/nni/algorithms/compression/v2/pytorch/test/test_pruner_torch.py
+++ b/nni/algorithms/compression/v2/pytorch/test/test_pruner_torch.py
@@ -14,8 +14,10 @@
     FPGMPruner,
     ActivationAPoZRankPruner,
     ActivationMeanRankPruner,
-    TaylorFOWeightPruner
+    TaylorFOWeightPruner,
+    ADMMPruner
 )
+from nni.algorithms.compression.v2.pytorch.utils.pruning import compute_sparsity_mask2compact
 
 
 class TorchModel(torch.nn.Module):
@@ -63,7 +65,9 @@ def test_level_pruner(self):
         config_list = [{'op_types': ['Conv2d'], 'sparsity': 0.8}]
         pruner = LevelPruner(model=model, config_list=config_list)
         pruned_model, masks = pruner.compress()
-        # TODO: validate masks sparsity
+        pruner._unwrap_model()
+        sparsity_list = compute_sparsity_mask2compact(pruned_model, masks, config_list)
+        assert 0.79 < sparsity_list[0]['total_sparsity'] < 0.81
 
     def test_l1_norm_pruner(self):
         model = TorchModel()
@@ -71,6 +75,9 @@ def test_l1_norm_pruner(self):
         pruner = L1NormPruner(model=model, config_list=config_list, mode='dependency_aware',
                               dummy_input=torch.rand(10, 1, 28, 28))
         pruned_model, masks = pruner.compress()
+        pruner._unwrap_model()
+        sparsity_list = compute_sparsity_mask2compact(pruned_model, masks, config_list)
+        assert 0.79 < sparsity_list[0]['total_sparsity'] < 0.81
 
     def test_l2_norm_pruner(self):
         model = TorchModel()
@@ -78,6 +85,9 @@ def test_l2_norm_pruner(self):
         pruner = L2NormPruner(model=model, config_list=config_list, mode='dependency_aware',
                               dummy_input=torch.rand(10, 1, 28, 28))
         pruned_model, masks = pruner.compress()
+        pruner._unwrap_model()
+        sparsity_list = compute_sparsity_mask2compact(pruned_model, masks, config_list)
+        assert 0.79 < sparsity_list[0]['total_sparsity'] < 0.81
 
     def test_fpgm_pruner(self):
         model = TorchModel()
@@ -85,6 +95,9 @@ def test_fpgm_pruner(self):
         pruner = FPGMPruner(model=model, config_list=config_list, mode='dependency_aware',
                             dummy_input=torch.rand(10, 1, 28, 28))
         pruned_model, masks = pruner.compress()
+        pruner._unwrap_model()
+        sparsity_list = compute_sparsity_mask2compact(pruned_model, masks, config_list)
+        assert 0.79 < sparsity_list[0]['total_sparsity'] < 0.81
 
     def test_slim_pruner(self):
         model = TorchModel()
@@ -92,6 +105,9 @@ def test_slim_pruner(self):
         pruner = SlimPruner(model=model, config_list=config_list, trainer=trainer, optimizer=get_optimizer(model),
                             criterion=criterion, training_epochs=1, scale=0.001, mode='global')
         pruned_model, masks = pruner.compress()
+        pruner._unwrap_model()
+        sparsity_list = compute_sparsity_mask2compact(pruned_model, masks, config_list)
+        assert 0.79 < sparsity_list[0]['total_sparsity'] < 0.81
 
     def test_activation_apoz_rank_pruner(self):
         model = TorchModel()
@@ -101,6 +117,9 @@ def test_activation_apoz_rank_pruner(self):
                                           activation='relu', mode='dependency_aware',
                                           dummy_input=torch.rand(10, 1, 28, 28))
         pruned_model, masks = pruner.compress()
+        pruner._unwrap_model()
+        sparsity_list = compute_sparsity_mask2compact(pruned_model, masks, config_list)
+        assert 0.79 < sparsity_list[0]['total_sparsity'] < 0.81
 
     def test_activation_mean_rank_pruner(self):
         model = TorchModel()
@@ -110,6 +129,9 @@ def test_activation_mean_rank_pruner(self):
                                           activation='relu', mode='dependency_aware',
                                           dummy_input=torch.rand(10, 1, 28, 28))
         pruned_model, masks = pruner.compress()
+        pruner._unwrap_model()
+        sparsity_list = compute_sparsity_mask2compact(pruned_model, masks, config_list)
+        assert 0.79 < sparsity_list[0]['total_sparsity'] < 0.81
 
     def test_taylor_fo_pruner(self):
         model = TorchModel()
@@ -118,7 +140,19 @@ def test_taylor_fo_pruner(self):
                                       optimizer=get_optimizer(model), criterion=criterion, training_batches=1,
                                       mode='dependency_aware', dummy_input=torch.rand(10, 1, 28, 28))
         pruned_model, masks = pruner.compress()
+        pruner._unwrap_model()
+        sparsity_list = compute_sparsity_mask2compact(pruned_model, masks, config_list)
+        assert 0.79 < sparsity_list[0]['total_sparsity'] < 0.81
 
+    def test_admm_pruner(self):
+        model = TorchModel()
+        config_list = [{'op_types': ['Conv2d'], 'sparsity': 0.8, 'rho': 1e-3}]
+        pruner = ADMMPruner(model=model, config_list=config_list, trainer=trainer, optimizer=get_optimizer(model),
+                            criterion=criterion, iterations=2, training_epochs=1)
+        pruned_model, masks = pruner.compress()
+        pruner._unwrap_model()
+        sparsity_list = compute_sparsity_mask2compact(pruned_model, masks, config_list)
+        assert 0.79 < sparsity_list[0]['total_sparsity'] < 0.81
 
 if __name__ == '__main__':
     unittest.main()

From 11a4c2c193117bd23f066bade0198cd3274d3d7b Mon Sep 17 00:00:00 2001
From: J-shang <shangning128@163.com>
Date: Sun, 26 Sep 2021 14:11:39 +0800
Subject: [PATCH 08/16] update example

---
 .../pruning/v2/scheduler_torch.py             | 22 +++++++++++++++++--
 .../v2/pytorch/pruning/tools/base.py          | 19 ++++++++--------
 2 files changed, 29 insertions(+), 12 deletions(-)

diff --git a/examples/model_compress/pruning/v2/scheduler_torch.py b/examples/model_compress/pruning/v2/scheduler_torch.py
index cb31db1f11..3d951df450 100644
--- a/examples/model_compress/pruning/v2/scheduler_torch.py
+++ b/examples/model_compress/pruning/v2/scheduler_torch.py
@@ -61,19 +61,37 @@ def evaluator(model):
     optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4)
     criterion = torch.nn.CrossEntropyLoss()
 
+    # pre-train the model
     for i in range(5):
         trainer(model, optimizer, criterion, i)
 
+    # the finetuner used in the scheduler should only have model as input
     finetuner = functools.partial(trainer,
                                   optimizer=torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4),
                                   criterion=criterion,
                                   epoch='PF')
 
     config_list = [{'op_types': ['Conv2d'], 'sparsity': 0.8}]
-    task_generator = AGPTaskGenerator(10, model, config_list, keep_intermidiate_result=True)
+
+    # Make sure initialize task generator at first, this because the model pass to the generator should be an unwrapped model.
+    # If you want to initialize pruner at first, you can use the follow code.
+
+    # pruner = L1NormPruner(model, config_list)
+    # pruner._unwrap_model()
+    # task_generator = AGPTaskGenerator(10, model, config_list, log_dir='.', keep_intermidiate_result=True)
+    # pruner._wrap_model()
+
+    # you can specify the log_dir, all intermidiate results and best result will save under this folder.
+    # if you don't want to keep intermidiate results, you can set `keep_intermidiate_result=False`.
+    task_generator = AGPTaskGenerator(10, model, config_list, log_dir='.', keep_intermidiate_result=True)
     pruner = L1NormPruner(model, config_list)
 
     dummy_input = torch.rand(10, 3, 32, 32).to(device)
-    scheduler = PruningScheduler(pruner, task_generator, finetuner=finetuner, speed_up=True, dummy_input=dummy_input, evaluator=evaluator)
+
+    # if you just want to keep the final result as the best result, you can pass evaluator as None.
+    # or the result with the highest score (given by evaluator) will be the best result.
+
+    # scheduler = PruningScheduler(pruner, task_generator, finetuner=finetuner, speed_up=True, dummy_input=dummy_input, evaluator=evaluator)
+    scheduler = PruningScheduler(pruner, task_generator, finetuner=finetuner, speed_up=True, dummy_input=dummy_input, evaluator=None)
 
     scheduler.compress()
diff --git a/nni/algorithms/compression/v2/pytorch/pruning/tools/base.py b/nni/algorithms/compression/v2/pytorch/pruning/tools/base.py
index 28d3627ba1..fed982bdc0 100644
--- a/nni/algorithms/compression/v2/pytorch/pruning/tools/base.py
+++ b/nni/algorithms/compression/v2/pytorch/pruning/tools/base.py
@@ -506,16 +506,15 @@ def _save_data(self, folder_name: str, model: Module, masks: Dict[str, Dict[str,
 
     def update_best_result(self, task_result: TaskResult):
         score = task_result.score
-        if score is not None:
-            task_id = task_result.task_id
-            task = self._tasks[task_id]
-            task.score = score
-            if self._best_score is None or score > self._best_score:
-                self._best_score = score
-                self._best_task_id = task_id
-                with Path(task.config_list_path).open('r') as fr:
-                    best_config_list = json_tricks.load(fr)
-                self._save_data('best_result', task_result.compact_model, task_result.compact_model_masks, best_config_list)
+        task_id = task_result.task_id
+        task = self._tasks[task_id]
+        task.score = score
+        if self._best_score is None or score > self._best_score:
+            self._best_score = score
+            self._best_task_id = task_id
+            with Path(task.config_list_path).open('r') as fr:
+                best_config_list = json_tricks.load(fr)
+            self._save_data('best_result', task_result.compact_model, task_result.compact_model_masks, best_config_list)
 
     def init_pending_tasks(self) -> List[Task]:
         raise NotImplementedError()

From 1961a54515f7e1e01a4210f5152b59de7150fb3b Mon Sep 17 00:00:00 2001
From: J-shang <shangning128@163.com>
Date: Sun, 26 Sep 2021 14:15:12 +0800
Subject: [PATCH 09/16] move test file

---
 .../test_pruner_torch.py => test/ut/sdk/test_v2_pruner_torch.py   | 0
 .../ut/sdk/test_v2_pruning_tools_torch.py                         | 0
 .../test/test_scheduler.py => test/ut/sdk/test_v2_scheduler.py    | 0
 .../ut/sdk/test_v2_task_generator.py                              | 0
 4 files changed, 0 insertions(+), 0 deletions(-)
 rename nni/algorithms/compression/v2/pytorch/test/test_pruner_torch.py => test/ut/sdk/test_v2_pruner_torch.py (100%)
 rename nni/algorithms/compression/v2/pytorch/test/test_pruning_tools_torch.py => test/ut/sdk/test_v2_pruning_tools_torch.py (100%)
 rename nni/algorithms/compression/v2/pytorch/test/test_scheduler.py => test/ut/sdk/test_v2_scheduler.py (100%)
 rename nni/algorithms/compression/v2/pytorch/test/test_task_generator.py => test/ut/sdk/test_v2_task_generator.py (100%)

diff --git a/nni/algorithms/compression/v2/pytorch/test/test_pruner_torch.py b/test/ut/sdk/test_v2_pruner_torch.py
similarity index 100%
rename from nni/algorithms/compression/v2/pytorch/test/test_pruner_torch.py
rename to test/ut/sdk/test_v2_pruner_torch.py
diff --git a/nni/algorithms/compression/v2/pytorch/test/test_pruning_tools_torch.py b/test/ut/sdk/test_v2_pruning_tools_torch.py
similarity index 100%
rename from nni/algorithms/compression/v2/pytorch/test/test_pruning_tools_torch.py
rename to test/ut/sdk/test_v2_pruning_tools_torch.py
diff --git a/nni/algorithms/compression/v2/pytorch/test/test_scheduler.py b/test/ut/sdk/test_v2_scheduler.py
similarity index 100%
rename from nni/algorithms/compression/v2/pytorch/test/test_scheduler.py
rename to test/ut/sdk/test_v2_scheduler.py
diff --git a/nni/algorithms/compression/v2/pytorch/test/test_task_generator.py b/test/ut/sdk/test_v2_task_generator.py
similarity index 100%
rename from nni/algorithms/compression/v2/pytorch/test/test_task_generator.py
rename to test/ut/sdk/test_v2_task_generator.py

From b79b20c4fbc9e753df0106687c66f20587366d3f Mon Sep 17 00:00:00 2001
From: J-shang <shangning128@163.com>
Date: Mon, 27 Sep 2021 09:35:22 +0800
Subject: [PATCH 10/16] fix import error

---
 nni/algorithms/compression/v2/pytorch/base/scheduler.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/nni/algorithms/compression/v2/pytorch/base/scheduler.py b/nni/algorithms/compression/v2/pytorch/base/scheduler.py
index 9923e93d64..8aafdcd116 100644
--- a/nni/algorithms/compression/v2/pytorch/base/scheduler.py
+++ b/nni/algorithms/compression/v2/pytorch/base/scheduler.py
@@ -9,8 +9,8 @@
 
 import json_tricks
 import torch
+from torch import Tensor
 from torch.nn import Module
-from torch.tensor import Tensor
 
 _logger = logging.getLogger(__name__)
 

From 041dd067a1b219319d036bb64415b6aeef7f8d85 Mon Sep 17 00:00:00 2001
From: J-shang <shangning128@163.com>
Date: Mon, 27 Sep 2021 10:17:02 +0800
Subject: [PATCH 11/16] adapt old python version

---
 nni/algorithms/compression/v2/pytorch/base/compressor.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/nni/algorithms/compression/v2/pytorch/base/compressor.py b/nni/algorithms/compression/v2/pytorch/base/compressor.py
index c2801d38f9..e0cfd0f1f1 100644
--- a/nni/algorithms/compression/v2/pytorch/base/compressor.py
+++ b/nni/algorithms/compression/v2/pytorch/base/compressor.py
@@ -3,7 +3,7 @@
 
 import collections
 import logging
-from typing import List, Dict, Optional, OrderedDict, Tuple, Any
+from typing import List, Dict, Optional, Tuple, Any
 
 import torch
 from torch.nn import Module
@@ -149,7 +149,7 @@ def _select_config(self, layer: LayerInfo) -> Optional[Dict]:
             return None
         return ret
 
-    def get_modules_wrapper(self) -> OrderedDict[str, Module]:
+    def get_modules_wrapper(self) -> Dict[str, Module]:
         """
         Returns
         -------

From 7b37ca820ee9456dbbc305be20408d5eaba61906 Mon Sep 17 00:00:00 2001
From: J-shang <shangning128@163.com>
Date: Tue, 28 Sep 2021 15:36:11 +0800
Subject: [PATCH 12/16] fix bug

---
 test/ut/sdk/test_v2_task_generator.py | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/test/ut/sdk/test_v2_task_generator.py b/test/ut/sdk/test_v2_task_generator.py
index 8e98308104..a248efc4e8 100644
--- a/test/ut/sdk/test_v2_task_generator.py
+++ b/test/ut/sdk/test_v2_task_generator.py
@@ -37,7 +37,7 @@ def forward(self, x):
         return F.log_softmax(x, dim=1)
 
 
-def test_task_generator(task_generator_type):
+def run_task_generator(task_generator_type):
     model = TorchModel()
     config_list = [{'op_types': ['Conv2d'], 'sparsity': 0.8}]
 
@@ -50,7 +50,7 @@ def test_task_generator(task_generator_type):
     elif task_generator_type == 'simulated_annealing':
         task_generator = SimulatedAnnealingTaskGenerator(model, config_list)
 
-    count = run_task_generator(task_generator)
+    count = run_task_generator_(task_generator)
 
     if task_generator_type == 'agp':
         assert count == 6
@@ -62,7 +62,7 @@ def test_task_generator(task_generator_type):
         assert count == 17
 
 
-def run_task_generator(task_generator):
+def run_task_generator_(task_generator):
     task = task_generator.next()
     factor = 0.9
     count = 0
@@ -77,16 +77,16 @@ def run_task_generator(task_generator):
 
 class TaskGenerator(unittest.TestCase):
     def test_agp_task_generator(self):
-        test_task_generator('agp')
+        run_task_generator('agp')
 
     def test_linear_task_generator(self):
-        test_task_generator('linear')
+        run_task_generator('linear')
 
     def test_lottery_ticket_task_generator(self):
-        test_task_generator('lottery_ticket')
+        run_task_generator('lottery_ticket')
 
     def test_simulated_annealing_task_generator(self):
-        test_task_generator('simulated_annealing')
+        run_task_generator('simulated_annealing')
 
 
 if __name__ == '__main__':

From d3df0b910e0f195a52d9905659f12a77b77d2c3d Mon Sep 17 00:00:00 2001
From: J-shang <shangning128@163.com>
Date: Tue, 28 Sep 2021 16:43:03 +0800
Subject: [PATCH 13/16] adapt old python version

---
 nni/algorithms/compression/v2/pytorch/base/scheduler.py       | 4 ++--
 .../v2/pytorch/pruning/tools/sparsity_allocator.py            | 2 ++
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/nni/algorithms/compression/v2/pytorch/base/scheduler.py b/nni/algorithms/compression/v2/pytorch/base/scheduler.py
index 8aafdcd116..0be80d7ee6 100644
--- a/nni/algorithms/compression/v2/pytorch/base/scheduler.py
+++ b/nni/algorithms/compression/v2/pytorch/base/scheduler.py
@@ -5,7 +5,7 @@
 import logging
 import os
 from pathlib import Path
-from typing import List, Dict, Tuple, Literal, Optional
+from typing import List, Dict, Tuple, Optional
 
 import json_tricks
 import torch
@@ -37,7 +37,7 @@ def __init__(self, task_id: int, model_path: str, masks_path: str, config_list_p
         self.masks_path = masks_path
         self.config_list_path = config_list_path
 
-        self.status: Literal['Pending', 'Running', 'Finished'] = 'Pending'
+        self.status = 'Pending'
         self.score: Optional[float] = None
 
         self.state = {}
diff --git a/nni/algorithms/compression/v2/pytorch/pruning/tools/sparsity_allocator.py b/nni/algorithms/compression/v2/pytorch/pruning/tools/sparsity_allocator.py
index f0f0e97961..ec84cfe130 100644
--- a/nni/algorithms/compression/v2/pytorch/pruning/tools/sparsity_allocator.py
+++ b/nni/algorithms/compression/v2/pytorch/pruning/tools/sparsity_allocator.py
@@ -103,8 +103,10 @@ def __init__(self, pruner: Pruner, dim: int, dummy_input: Any):
 
     def _get_dependency(self):
         graph = self.pruner.generate_graph(dummy_input=self.dummy_input)
+        self.pruner._unwrap_model()
         self.channel_depen = ChannelDependency(model=self.pruner.bound_model, dummy_input=self.dummy_input, traced_model=graph.trace).dependency_sets
         self.group_depen = GroupDependency(model=self.pruner.bound_model, dummy_input=self.dummy_input, traced_model=graph.trace).dependency_sets
+        self.pruner._wrap_model()
 
     def generate_sparsity(self, metrics: Dict) -> Dict[str, Dict[str, Tensor]]:
         self._get_dependency()

From 626b08200a6b50c83d3be43d44b17efeb3a824f5 Mon Sep 17 00:00:00 2001
From: J-shang <shangning128@163.com>
Date: Wed, 29 Sep 2021 22:02:03 +0800
Subject: [PATCH 14/16] update example

---
 .../pruning/v2/scheduler_torch.py              | 18 ++++++++++++------
 1 file changed, 12 insertions(+), 6 deletions(-)

diff --git a/examples/model_compress/pruning/v2/scheduler_torch.py b/examples/model_compress/pruning/v2/scheduler_torch.py
index 3d951df450..2d3a873627 100644
--- a/examples/model_compress/pruning/v2/scheduler_torch.py
+++ b/examples/model_compress/pruning/v2/scheduler_torch.py
@@ -42,6 +42,18 @@ def trainer(model, optimizer, criterion, epoch):
         loss.backward()
         optimizer.step()
 
+def finetuner(model):
+    model.train()
+    optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4)
+    criterion = torch.nn.CrossEntropyLoss()
+    for data, target in tqdm(iterable=train_loader, desc='Epoch PFs'):
+        data, target = data.to(device), target.to(device)
+        optimizer.zero_grad()
+        output = model(data)
+        loss = criterion(output, target)
+        loss.backward()
+        optimizer.step()
+
 def evaluator(model):
     model.eval()
     correct = 0
@@ -65,12 +77,6 @@ def evaluator(model):
     for i in range(5):
         trainer(model, optimizer, criterion, i)
 
-    # the finetuner used in the scheduler should only have model as input
-    finetuner = functools.partial(trainer,
-                                  optimizer=torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4),
-                                  criterion=criterion,
-                                  epoch='PF')
-
     config_list = [{'op_types': ['Conv2d'], 'sparsity': 0.8}]
 
     # Make sure initialize task generator at first, this because the model pass to the generator should be an unwrapped model.

From 6effafa43a4d5e7cdb61ede1955adc2542a0ecb7 Mon Sep 17 00:00:00 2001
From: J-shang <shangning128@163.com>
Date: Fri, 8 Oct 2021 14:41:22 +0800
Subject: [PATCH 15/16] add __init__ in utils

---
 .../compression/v2/pytorch/base/compressor.py         |  2 +-
 .../compression/v2/pytorch/pruning/basic_pruner.py    |  3 +--
 .../v2/pytorch/pruning/tools/task_generator.py        |  2 +-
 .../compression/v2/pytorch/utils/__init__.py          | 11 +++++++++++
 test/ut/sdk/test_v2_pruner_torch.py                   |  2 +-
 test/ut/sdk/test_v2_pruning_tools_torch.py            |  2 +-
 6 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/nni/algorithms/compression/v2/pytorch/base/compressor.py b/nni/algorithms/compression/v2/pytorch/base/compressor.py
index e0cfd0f1f1..0e288fa092 100644
--- a/nni/algorithms/compression/v2/pytorch/base/compressor.py
+++ b/nni/algorithms/compression/v2/pytorch/base/compressor.py
@@ -9,7 +9,7 @@
 from torch.nn import Module
 
 from nni.common.graph_utils import TorchModuleGraph
-from nni.algorithms.compression.v2.pytorch.utils.pruning import get_module_by_name
+from nni.algorithms.compression.v2.pytorch.utils import get_module_by_name
 
 _logger = logging.getLogger(__name__)
 
diff --git a/nni/algorithms/compression/v2/pytorch/pruning/basic_pruner.py b/nni/algorithms/compression/v2/pytorch/pruning/basic_pruner.py
index 8157b28ce7..c9ebcdc797 100644
--- a/nni/algorithms/compression/v2/pytorch/pruning/basic_pruner.py
+++ b/nni/algorithms/compression/v2/pytorch/pruning/basic_pruner.py
@@ -13,8 +13,7 @@
 from torch.optim import Optimizer
 
 from nni.algorithms.compression.v2.pytorch.base.pruner import Pruner
-from nni.algorithms.compression.v2.pytorch.utils.config_validation import CompressorSchema
-from nni.algorithms.compression.v2.pytorch.utils.pruning import config_list_canonical
+from nni.algorithms.compression.v2.pytorch.utils import CompressorSchema, config_list_canonical
 
 from .tools import (
     DataCollector,
diff --git a/nni/algorithms/compression/v2/pytorch/pruning/tools/task_generator.py b/nni/algorithms/compression/v2/pytorch/pruning/tools/task_generator.py
index c6133869b3..0a3a09f65b 100644
--- a/nni/algorithms/compression/v2/pytorch/pruning/tools/task_generator.py
+++ b/nni/algorithms/compression/v2/pytorch/pruning/tools/task_generator.py
@@ -13,7 +13,7 @@
 from torch.nn import Module
 
 from nni.algorithms.compression.v2.pytorch.base import Task, TaskResult
-from nni.algorithms.compression.v2.pytorch.utils.pruning import (
+from nni.algorithms.compression.v2.pytorch.utils import (
     config_list_canonical,
     compute_sparsity,
     get_model_weights_numel
diff --git a/nni/algorithms/compression/v2/pytorch/utils/__init__.py b/nni/algorithms/compression/v2/pytorch/utils/__init__.py
index e69de29bb2..1777c499c9 100644
--- a/nni/algorithms/compression/v2/pytorch/utils/__init__.py
+++ b/nni/algorithms/compression/v2/pytorch/utils/__init__.py
@@ -0,0 +1,11 @@
+from .config_validation import CompressorSchema
+from .pruning import (
+    config_list_canonical,
+    unfold_config_list,
+    dedupe_config_list,
+    compute_sparsity_compact2origin,
+    compute_sparsity_mask2compact,
+    compute_sparsity,
+    get_model_weights_numel,
+    get_module_by_name
+)
diff --git a/test/ut/sdk/test_v2_pruner_torch.py b/test/ut/sdk/test_v2_pruner_torch.py
index a727ecaeb8..50b3c22fe2 100644
--- a/test/ut/sdk/test_v2_pruner_torch.py
+++ b/test/ut/sdk/test_v2_pruner_torch.py
@@ -17,7 +17,7 @@
     TaylorFOWeightPruner,
     ADMMPruner
 )
-from nni.algorithms.compression.v2.pytorch.utils.pruning import compute_sparsity_mask2compact
+from nni.algorithms.compression.v2.pytorch.utils import compute_sparsity_mask2compact
 
 
 class TorchModel(torch.nn.Module):
diff --git a/test/ut/sdk/test_v2_pruning_tools_torch.py b/test/ut/sdk/test_v2_pruning_tools_torch.py
index 579dd6f678..625d380bb4 100644
--- a/test/ut/sdk/test_v2_pruning_tools_torch.py
+++ b/test/ut/sdk/test_v2_pruning_tools_torch.py
@@ -24,7 +24,7 @@
     GlobalSparsityAllocator
 )
 from nni.algorithms.compression.v2.pytorch.pruning.tools.base import HookCollectorInfo
-from nni.algorithms.compression.v2.pytorch.utils.pruning import get_module_by_name
+from nni.algorithms.compression.v2.pytorch.utils import get_module_by_name
 
 
 class TorchModel(torch.nn.Module):

From 6d4a46ee25c5d2300268020e0d3d74fc1609239d Mon Sep 17 00:00:00 2001
From: J-shang <shangning128@163.com>
Date: Tue, 12 Oct 2021 09:23:44 +0800
Subject: [PATCH 16/16] fix spelling

---
 .../pruning/v2/scheduler_torch.py             |  8 ++---
 .../v2/pytorch/pruning/tools/base.py          | 12 ++++----
 .../pytorch/pruning/tools/task_generator.py   | 30 +++++++++----------
 3 files changed, 25 insertions(+), 25 deletions(-)

diff --git a/examples/model_compress/pruning/v2/scheduler_torch.py b/examples/model_compress/pruning/v2/scheduler_torch.py
index 2d3a873627..d020bd131c 100644
--- a/examples/model_compress/pruning/v2/scheduler_torch.py
+++ b/examples/model_compress/pruning/v2/scheduler_torch.py
@@ -84,12 +84,12 @@ def evaluator(model):
 
     # pruner = L1NormPruner(model, config_list)
     # pruner._unwrap_model()
-    # task_generator = AGPTaskGenerator(10, model, config_list, log_dir='.', keep_intermidiate_result=True)
+    # task_generator = AGPTaskGenerator(10, model, config_list, log_dir='.', keep_intermediate_result=True)
     # pruner._wrap_model()
 
-    # you can specify the log_dir, all intermidiate results and best result will save under this folder.
-    # if you don't want to keep intermidiate results, you can set `keep_intermidiate_result=False`.
-    task_generator = AGPTaskGenerator(10, model, config_list, log_dir='.', keep_intermidiate_result=True)
+    # you can specify the log_dir, all intermediate results and best result will save under this folder.
+    # if you don't want to keep intermediate results, you can set `keep_intermediate_result=False`.
+    task_generator = AGPTaskGenerator(10, model, config_list, log_dir='.', keep_intermediate_result=True)
     pruner = L1NormPruner(model, config_list)
 
     dummy_input = torch.rand(10, 3, 32, 32).to(device)
diff --git a/nni/algorithms/compression/v2/pytorch/pruning/tools/base.py b/nni/algorithms/compression/v2/pytorch/pruning/tools/base.py
index fed982bdc0..ed776472c5 100644
--- a/nni/algorithms/compression/v2/pytorch/pruning/tools/base.py
+++ b/nni/algorithms/compression/v2/pytorch/pruning/tools/base.py
@@ -452,7 +452,7 @@ class TaskGenerator:
     This class used to generate config list for pruner in each iteration.
     """
     def __init__(self, origin_model: Module, origin_masks: Dict[str, Dict[str, Tensor]] = {},
-                 origin_config_list: List[Dict] = [], log_dir: str = '.', keep_intermidiate_result: bool = False):
+                 origin_config_list: List[Dict] = [], log_dir: str = '.', keep_intermediate_result: bool = False):
         """
         Parameters
         ----------
@@ -465,16 +465,16 @@ def __init__(self, origin_model: Module, origin_masks: Dict[str, Dict[str, Tenso
             This means the sparsity provided by the origin_masks should also be recorded in the origin_config_list.
         log_dir
             The log directory use to saving the task generator log.
-        keep_intermidiate_result
+        keep_intermediate_result
             If keeping the intermediate result, including intermediate model and masks during each iteration.
         """
         assert isinstance(origin_model, Module), 'Only support pytorch module.'
 
         self._log_dir_root = Path(log_dir, datetime.now().strftime('%Y-%m-%d-%H-%M-%S-%f')).absolute()
         self._log_dir_root.mkdir(parents=True, exist_ok=True)
-        self._keep_intermidiate_result = keep_intermidiate_result
-        self._intermidiate_result_dir = Path(self._log_dir_root, 'intermidiate_result')
-        self._intermidiate_result_dir.mkdir(parents=True, exist_ok=True)
+        self._keep_intermediate_result = keep_intermediate_result
+        self._intermediate_result_dir = Path(self._log_dir_root, 'intermediate_result')
+        self._intermediate_result_dir.mkdir(parents=True, exist_ok=True)
 
         # save origin data in {log_dir}/origin
         self._origin_model_path = Path(self._log_dir_root, 'origin', 'model.pth')
@@ -539,7 +539,7 @@ def receive_task_result(self, task_result: TaskResult):
         self._pending_tasks.extend(self.generate_tasks(task_result))
         self._dump_tasks_info()
 
-        if not self._keep_intermidiate_result:
+        if not self._keep_intermediate_result:
             self._tasks[task_id].clean_up()
 
     def next(self) -> Optional[Task]:
diff --git a/nni/algorithms/compression/v2/pytorch/pruning/tools/task_generator.py b/nni/algorithms/compression/v2/pytorch/pruning/tools/task_generator.py
index 0a3a09f65b..e141bf4ffa 100644
--- a/nni/algorithms/compression/v2/pytorch/pruning/tools/task_generator.py
+++ b/nni/algorithms/compression/v2/pytorch/pruning/tools/task_generator.py
@@ -25,7 +25,7 @@
 
 class FunctionBasedTaskGenerator(TaskGenerator):
     def __init__(self, total_iteration: int, origin_model: Module, origin_config_list: List[Dict],
-                 origin_masks: Dict[str, Dict[str, Tensor]] = {}, log_dir: str = '.', keep_intermidiate_result: bool = False):
+                 origin_masks: Dict[str, Dict[str, Tensor]] = {}, log_dir: str = '.', keep_intermediate_result: bool = False):
         """
         Parameters
         ----------
@@ -40,7 +40,7 @@ def __init__(self, total_iteration: int, origin_model: Module, origin_config_lis
             The pre masks on the origin model. This mask maybe user-defined or maybe generate by previous pruning.
         log_dir
             The log directory use to saving the task generator log.
-        keep_intermidiate_result
+        keep_intermediate_result
             If keeping the intermediate result, including intermediate model and masks during each iteration.
         """
         self.current_iteration = 0
@@ -48,7 +48,7 @@ def __init__(self, total_iteration: int, origin_model: Module, origin_config_lis
         self.total_iteration = total_iteration
 
         super().__init__(origin_model, origin_config_list=self.target_sparsity, origin_masks=origin_masks,
-                         log_dir=log_dir, keep_intermidiate_result=keep_intermidiate_result)
+                         log_dir=log_dir, keep_intermediate_result=keep_intermediate_result)
 
     def init_pending_tasks(self) -> List[Task]:
         origin_model = torch.load(self._origin_model_path)
@@ -62,9 +62,9 @@ def generate_tasks(self, task_result: TaskResult) -> List[Task]:
         compact_model = task_result.compact_model
         compact_model_masks = task_result.compact_model_masks
 
-        # save intermidiate result
-        model_path = Path(self._intermidiate_result_dir, '{}_compact_model.pth'.format(task_result.task_id))
-        masks_path = Path(self._intermidiate_result_dir, '{}_compact_model_masks.pth'.format(task_result.task_id))
+        # save intermediate result
+        model_path = Path(self._intermediate_result_dir, '{}_compact_model.pth'.format(task_result.task_id))
+        masks_path = Path(self._intermediate_result_dir, '{}_compact_model_masks.pth'.format(task_result.task_id))
         torch.save(compact_model, model_path)
         torch.save(compact_model_masks, masks_path)
 
@@ -81,7 +81,7 @@ def generate_tasks(self, task_result: TaskResult) -> List[Task]:
 
         task_id = self._task_id_candidate
         new_config_list = self.generate_config_list(self.target_sparsity, self.current_iteration, compact2origin_sparsity)
-        config_list_path = Path(self._intermidiate_result_dir, '{}_config_list.json'.format(task_id))
+        config_list_path = Path(self._intermediate_result_dir, '{}_config_list.json'.format(task_id))
 
         with Path(config_list_path).open('w') as f:
             json_tricks.dump(new_config_list, f, indent=4)
@@ -124,9 +124,9 @@ def generate_config_list(self, target_sparsity: List[Dict], iteration: int, comp
 
 class LotteryTicketTaskGenerator(FunctionBasedTaskGenerator):
     def __init__(self, total_iteration: int, origin_model: Module, origin_config_list: List[Dict],
-                 origin_masks: Dict[str, Dict[str, Tensor]] = {}, log_dir: str = '.', keep_intermidiate_result: bool = False):
+                 origin_masks: Dict[str, Dict[str, Tensor]] = {}, log_dir: str = '.', keep_intermediate_result: bool = False):
         super().__init__(total_iteration, origin_model, origin_config_list, origin_masks=origin_masks, log_dir=log_dir,
-                         keep_intermidiate_result=keep_intermidiate_result)
+                         keep_intermediate_result=keep_intermediate_result)
         self.current_iteration = 1
 
     def generate_config_list(self, target_sparsity: List[Dict], iteration: int, compact2origin_sparsity: List[Dict]) -> List[Dict]:
@@ -147,7 +147,7 @@ def generate_config_list(self, target_sparsity: List[Dict], iteration: int, comp
 class SimulatedAnnealingTaskGenerator(TaskGenerator):
     def __init__(self, origin_model: Module, origin_config_list: List[Dict], origin_masks: Dict[str, Dict[str, Tensor]] = {},
                  start_temperature: float = 100, stop_temperature: float = 20, cool_down_rate: float = 0.9,
-                 perturbation_magnitude: float = 0.35, log_dir: str = '.', keep_intermidiate_result: bool = False):
+                 perturbation_magnitude: float = 0.35, log_dir: str = '.', keep_intermediate_result: bool = False):
         """
         Parameters
         ----------
@@ -168,7 +168,7 @@ def __init__(self, origin_model: Module, origin_config_list: List[Dict], origin_
             Initial perturbation magnitude to the sparsities. The magnitude decreases with current temperature.
         log_dir
             The log directory use to saving the task generator log.
-        keep_intermidiate_result
+        keep_intermediate_result
             If keeping the intermediate result, including intermediate model and masks during each iteration.
         """
         self.start_temperature = start_temperature
@@ -186,7 +186,7 @@ def __init__(self, origin_model: Module, origin_config_list: List[Dict], origin_
         self._current_score = None
 
         super().__init__(origin_model, origin_masks=origin_masks, origin_config_list=origin_config_list,
-                         log_dir=log_dir, keep_intermidiate_result=keep_intermidiate_result)
+                         log_dir=log_dir, keep_intermediate_result=keep_intermediate_result)
 
     def _adjust_target_sparsity(self):
         """
@@ -288,8 +288,8 @@ def init_pending_tasks(self) -> List[Task]:
         origin_model = torch.load(self._origin_model_path)
         origin_masks = torch.load(self._origin_masks_path)
 
-        self.temp_model_path = Path(self._intermidiate_result_dir, 'origin_compact_model.pth')
-        self.temp_masks_path = Path(self._intermidiate_result_dir, 'origin_compact_model_masks.pth')
+        self.temp_model_path = Path(self._intermediate_result_dir, 'origin_compact_model.pth')
+        self.temp_masks_path = Path(self._intermediate_result_dir, 'origin_compact_model_masks.pth')
         torch.save(origin_model, self.temp_model_path)
         torch.save(origin_masks, self.temp_masks_path)
 
@@ -319,7 +319,7 @@ def generate_tasks(self, task_result: TaskResult) -> List[Task]:
 
         task_id = self._task_id_candidate
         new_config_list = self._recover_real_sparsity(deepcopy(self._temp_config_list))
-        config_list_path = Path(self._intermidiate_result_dir, '{}_config_list.json'.format(task_id))
+        config_list_path = Path(self._intermediate_result_dir, '{}_config_list.json'.format(task_id))
 
         with Path(config_list_path).open('w') as f:
             json_tricks.dump(new_config_list, f, indent=4)