Repository: microsoft/nni
File: nni/algorithms/compression/v2/pytorch/pruning/basic_pruner.py (Python)

Code before:

    def _calculate_taylor_expansion(self, weight_tensor: Tensor, grad: Tensor) -> Tensor:
        return (weight_tensor.detach() * grad.detach()).data.pow(2)

    def reset_tools(self):
        hook_targets = {layer_info.name: layer_info.module.weight for layer_info, _ in self._detect_modules_to_compress()}
        collector_info = HookCollectorInfo(hook_targets, 'tensor', self._collector)
        if self.data_collector is None:
            self.data_collector = SingleHookTrainerBasedDataCollector(self, self.trainer, self.optimizer, self.criterion,
                                                                      1, collector_infos=[collector_info])
        else:
            self.data_collector.reset()
        if self.metrics_calculator is None:
            self.metrics_calculator = MultiDataNormMetricsCalculator(p=1, dim=0)
        if self.sparsity_allocator is None:
            if self.mode == 'normal':
                self.sparsity_allocator = NormalSparsityAllocator(self, dim=0)
            elif self.mode == 'global':
                self.sparsity_allocator = GlobalSparsityAllocator(self, dim=0)
            elif self.mode == 'dependency_aware':
                self.sparsity_allocator = Conv2dDependencyAwareAllocator(self, 0, self.dummy_input)
            else:
                raise NotImplementedError('Only support mode `normal`, `global` and `dependency_aware`')


class ADMMPruner(OneShotPruner):
    """
    ADMM (Alternating Direction Method of Multipliers) Pruner is a kind of mathematical optimization technique.
    The metric used in this pruner is the absolute value of the weight.
    In each iteration, the weight with small magnitudes will be set to zero.
    Only in the final iteration, the mask will be generated and apply to model wrapper.

    The original paper refer to: https://arxiv.org/abs/1804.03294.
    """

    def __init__(self, model: Module, config_list: List[Dict], trainer: Callable[[Module, Optimizer, Callable], None],
                 optimizer: Optimizer, criterion: Callable[[Tensor, Tensor], Tensor], iterations: int, training_epochs: int):
        """
        Parameters
        ----------
        model
            Model to be pruned.
        config_list
            Supported keys:
                - sparsity : This is to specify the sparsity for each layer in this config to be compressed.
                - sparsity_per_layer : Equals to sparsity.
                - row : Penalty parameters.
                - 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.
        trainer

Review comment (refactor):
It seems that ADMMPruner has some functionalities of the PruningScheduler, e.g., it performs multiple iterations and keeps track of context/state data like Z and U (I'm not sure whether this is supposed to be handled by the "task" abstraction we discussed). How do we want to integrate ADMMPruner with the scheduler logic?

Code after:

    def _calculate_taylor_expansion(self, weight_tensor: Tensor, grad: Tensor) -> Tensor:
        return (weight_tensor.detach() * grad.detach()).data.pow(2)

    def reset_tools(self):
        hook_targets = {layer_info.name: layer_info.module.weight for layer_info, _ in self._detect_modules_to_compress()}
        collector_info = HookCollectorInfo(hook_targets, 'tensor', self._collector)
        if self.data_collector is None:
            self.data_collector = SingleHookTrainerBasedDataCollector(self, self.trainer, self.optimizer, self.criterion,
                                                                      1, collector_infos=[collector_info])
        else:
            self.data_collector.reset()
        if self.metrics_calculator is None:
            self.metrics_calculator = MultiDataNormMetricsCalculator(p=1, dim=0)
        if self.sparsity_allocator is None:
            if self.mode == 'normal':
                self.sparsity_allocator = NormalSparsityAllocator(self, dim=0)
            elif self.mode == 'global':
                self.sparsity_allocator = GlobalSparsityAllocator(self, dim=0)
            elif self.mode == 'dependency_aware':
                self.sparsity_allocator = Conv2dDependencyAwareAllocator(self, 0, self.dummy_input)
            else:
                raise NotImplementedError('Only support mode `normal`, `global` and `dependency_aware`')


class ADMMPruner(BasicPruner):
    """
    ADMM (Alternating Direction Method of Multipliers) Pruner is a kind of mathematical optimization technique.
    The metric used in this pruner is the absolute value of the weight.
    In each iteration, the weight with small magnitudes will be set to zero.
    Only in the final iteration, the mask will be generated and apply to model wrapper.

    The original paper refer to: https://arxiv.org/abs/1804.03294.
    """

    def __init__(self, model: Module, config_list: List[Dict], trainer: Callable[[Module, Optimizer, Callable], None],
                 optimizer: Optimizer, criterion: Callable[[Tensor, Tensor], Tensor], iterations: int, training_epochs: int):
        """
        Parameters
        ----------
        model
            Model to be pruned.
        config_list
            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.
                - 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.
        trainer