[Model Compression] Add global sort for taylor pruner

docs/en_US/Compression/Pruner.rst

@@ -334,6 +334,8 @@ TaylorFOWeightFilter Pruner is a pruner which prunes convolutional layers based
 
 We also provide a dependency-aware mode for this pruner to get better speedup from the pruning. Please reference `dependency-aware <./DependencyAware.rst>`__ for more details.
 
+What's more, we provide a global-sort mode for this pruner which is aligned with paper implementation. Please set parameter 'global_sort' to True when instantiate TaylorFOWeightFilterPruner.
+
 Usage
 ^^^^^
 

examples/model_compress/pruning/basic_pruners_torch.py

@@ -218,6 +218,10 @@ def trainer(model, optimizer, criterion, epoch):
         }]
 
     else:
+        if args.global_sort:
+            print('Enable the global_sort mode')
+            # only taylor pruner supports global sort mode currently
+            kw_args['global_sort'] = True
         if args.dependency_aware:
             dummy_input = get_dummy_input(args, device)
             print('Enable the dependency_aware mode')
@@ -331,6 +335,8 @@ def trainer(model, optimizer, criterion, epoch):
                         help='target overall target sparsity')
     parser.add_argument('--dependency-aware', action='store_true', default=False,
                         help='toggle dependency aware mode')
+    parser.add_argument('--global-sort', action='store_true', default=False,
+                        help='toggle global sort mode')
     parser.add_argument('--pruner', type=str, default='l1filter',
                         choices=['level', 'l1filter', 'l2filter', 'slim', 'agp',
                                  'fpgm', 'mean_activation', 'apoz', 'taylorfo'],
@@ -356,4 +362,4 @@ def trainer(model, optimizer, criterion, epoch):
         args.pruner = params['pruner']
         args.model = params['model']
 
-    main(args)
+    main(args)

nni/algorithms/compression/pytorch/pruning/dependency_aware_pruner.py

@@ -27,7 +27,7 @@ class DependencyAwarePruner(Pruner):
     """
 
     def __init__(self, model, config_list, optimizer=None, pruning_algorithm='level', dependency_aware=False,
-                 dummy_input=None, **algo_kwargs):
+                 dummy_input=None, global_sort=False, **algo_kwargs):
         super().__init__(model, config_list=config_list, optimizer=optimizer)
 
         self.dependency_aware = dependency_aware
@@ -56,6 +56,7 @@ def __init__(self, model, config_list, optimizer=None, pruning_algorithm='level'
             model, self, **algo_kwargs)
         # set the dependency-aware switch for the masker
         self.masker.dependency_aware = dependency_aware
+        self.masker.global_sort = global_sort
         self.set_wrappers_attribute("if_calculated", False)
 
     def calc_mask(self, wrapper, wrapper_idx=None):

nni/algorithms/compression/pytorch/pruning/iterative_pruner.py

@@ -22,7 +22,7 @@ class IterativePruner(DependencyAwarePruner):
     """
 
     def __init__(self, model, config_list, optimizer=None, pruning_algorithm='slim', trainer=None, criterion=None,
-                 num_iterations=20, epochs_per_iteration=5, dependency_aware=False, dummy_input=None, **algo_kwargs):
+                 num_iterations=20, epochs_per_iteration=5, dependency_aware=False, dummy_input=None, global_sort=False, **algo_kwargs):
         """
         Parameters
         ----------
@@ -51,6 +51,9 @@ def __init__(self, model, config_list, optimizer=None, pruning_algorithm='slim',
         dummy_input: torch.Tensor
             The dummy input to analyze the topology constraints. Note that,
             the dummy_input should on the same device with the model.
+        global_sort: bool
+            If prune the model in a global-sort way.
+            Only support TaylorFOWeightFilterPruner currently.
         algo_kwargs: dict
             Additional parameters passed to pruning algorithm masker class
         """
@@ -486,10 +489,15 @@ class TaylorFOWeightFilterPruner(IterativePruner):
     dummy_input : torch.Tensor
         The dummy input to analyze the topology constraints. Note that, the dummy_input
         should on the same device with the model.
+    global_sort: bool
+        Only support TaylorFOWeightFilterPruner currently.
+        If prune the model in a global-sort way. If it is `True`, this pruner will prune
+        the model according to the global contributions information which means channel contributions
+        will be sorted globally and whether specific channel will be pruned depends on global information.
     """
 
     def __init__(self, model, config_list, optimizer, trainer, criterion, sparsifying_training_batches=1,
-                 dependency_aware=False, dummy_input=None):
+                 dependency_aware=False, dummy_input=None, global_sort=False):
         super().__init__(model, config_list, optimizer=optimizer, pruning_algorithm='taylorfo', trainer=trainer,
                          criterion=criterion, statistics_batch_num=sparsifying_training_batches, num_iterations=1,
                          epochs_per_iteration=1, dependency_aware=dependency_aware,

nni/algorithms/compression/pytorch/pruning/structured_pruning_masker.py

@@ -33,11 +33,12 @@ class StructuredWeightMasker(WeightMasker):
 
     """
 
-    def __init__(self, model, pruner, preserve_round=1, dependency_aware=False):
+    def __init__(self, model, pruner, preserve_round=1, dependency_aware=False, global_sort=False):
         self.model = model
         self.pruner = pruner
         self.preserve_round = preserve_round
         self.dependency_aware = dependency_aware
+        self.global_sort = global_sort
 
     def calc_mask(self, sparsity, wrapper, wrapper_idx=None, **depen_kwargs):
         """
@@ -60,7 +61,11 @@ def calc_mask(self, sparsity, wrapper, wrapper_idx=None, **depen_kwargs):
         depen_kwargs: dict
             The kw_args for the dependency-aware mode.
         """
-        if not self.dependency_aware:
+        if self.global_sort:
+            # if the global_sort switch is on, calculate the mask based
+            # on global model information 
+            return self._global_calc_mask(sparsity, wrapper, wrapper_idx)
+        elif not self.dependency_aware:
             # calculate the mask in the normal way, each layer calculate its
             # own mask separately
             return self._normal_calc_mask(sparsity, wrapper, wrapper_idx)
@@ -127,6 +132,12 @@ def _get_current_state(self, sparsity, wrapper, wrapper_idx=None):
         # weight*mask_weight: apply base mask for iterative pruning
         return mask, weight * mask_weight, num_prune
 
+    def _global_calc_mask(self, sparsity, wrapper, wrapper_idx=None):
+        num_prune = self._get_global_num_prune(wrapper, wrapper_idx)
+        mask, weight, _ = self._get_current_state(
+            sparsity, wrapper, wrapper_idx)
+        return self.get_mask(mask, weight, num_prune, wrapper, wrapper_idx)
+
     def _normal_calc_mask(self, sparsity, wrapper, wrapper_idx=None):
         """
         Calculate the mask of given layer.
@@ -477,6 +488,29 @@ def __init__(self, model, pruner, statistics_batch_num=1):
         self.pruner.iterations = 0
         self.pruner.set_wrappers_attribute("contribution", None)
         self.pruner.patch_optimizer(self.calc_contributions)
+        self.global_threshold = None
+
+    def _get_global_threshold(self):
+        channel_contribution_list = []
+        for wrapper_idx, wrapper in enumerate(self.pruner.get_modules_wrapper()):
+            channel_contribution = self.get_channel_sum(wrapper, wrapper_idx)
+            channel_contribution_list.append(channel_contribution)
+        all_channel_contributions = torch.cat(channel_contribution_list)
+        k = int(all_channel_contributions.shape[0] * self.pruner.config_list[0]['sparsity'])
+        self.global_threshold = torch.topk(
+            all_channel_contributions.view(-1), k, largest=False)[0].max()
+        print(f'set global threshold to {self.global_threshold}')
+
+    def _get_global_num_prune(self, wrapper, wrapper_idx):
+        if self.global_threshold is None:
+            self._get_global_threshold()
+        weight = wrapper.module.weight.data
+        filters = weight.size(0)
+        channel_contribution = self.get_channel_sum(wrapper, wrapper_idx)
+        num_prune = channel_contribution[channel_contribution < self.global_threshold].size()[0]
+        if num_prune == filters:
+            num_prune -= 1
+        return num_prune
 
     def get_mask(self, base_mask, weight, num_prune, wrapper, wrapper_idx, channel_masks=None):
         channel_contribution = self.get_channel_sum(wrapper, wrapper_idx)