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

Code before:
        config_list: list
            List on pruning configs
        optimizer: torch.optim.Optimizer
            Optimizer used to train model
        pruning_algorithm: str
            algorithms being used to prune model
        trainer: function
            Function used to train the model.
            Users should write this function as a normal function to train the Pytorch model
            and include `model, optimizer, criterion, epoch` as function arguments.
        criterion: function
            Function used to calculate the loss between the target and the output.
            For example, you can use ``torch.nn.CrossEntropyLoss()`` as input.
        num_iterations: int
            Total number of iterations in pruning process. We will calculate mask at the end of an iteration.
        epochs_per_iteration: Union[int, list]
            The number of training epochs for each iteration. `int` represents the same value for each iteration.
            `list` represents the specific value for each iteration.
        dependency_aware: bool
            If prune the model in a dependency-aware way.
        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
        """
        super().__init__(model, config_list, optimizer, pruning_algorithm, dependency_aware, dummy_input, global_sort, **algo_kwargs)

        if isinstance(epochs_per_iteration, list):
            assert len(epochs_per_iteration) == num_iterations, 'num_iterations should equal to the length of epochs_per_iteration'
            self.epochs_per_iteration = epochs_per_iteration
        else:
            assert num_iterations > 0, 'num_iterations should >= 1'
            self.epochs_per_iteration = [epochs_per_iteration] * num_iterations

        self._validate_iteration_params()

        self._trainer = trainer
        self._criterion = criterion

    def _fresh_calculated(self):
        for wrapper in self.get_modules_wrapper():
            wrapper.if_calculated = False

    def _validate_iteration_params(self):
        assert all(num >= 0 for num in self.epochs_per_iteration), 'all epoch number need >= 0'

    def compress(self):

Review comment (suggestion):
if only `TaylorFOWeightFilterPruner` supports `global_sort`, suggest to put this argument in `TaylorFOWeightFilterPruner`.

Code after:
        config_list: list
            List on pruning configs
        optimizer: torch.optim.Optimizer
            Optimizer used to train model
        pruning_algorithm: str
            algorithms being used to prune model
        trainer: function
            Function used to train the model.
            Users should write this function as a normal function to train the Pytorch model
            and include `model, optimizer, criterion, epoch` as function arguments.
        criterion: function
            Function used to calculate the loss between the target and the output.
            For example, you can use ``torch.nn.CrossEntropyLoss()`` as input.
        num_iterations: int
            Total number of iterations in pruning process. We will calculate mask at the end of an iteration.
        epochs_per_iteration: Union[int, list]
            The number of training epochs for each iteration. `int` represents the same value for each iteration.
            `list` represents the specific value for each iteration.
        dependency_aware: bool
            If prune the model in a dependency-aware way.
        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.
        algo_kwargs: dict
            Additional parameters passed to pruning algorithm masker class
        """
        super().__init__(model, config_list, optimizer, pruning_algorithm, dependency_aware, dummy_input, **algo_kwargs)

        if isinstance(epochs_per_iteration, list):
            assert len(epochs_per_iteration) == num_iterations, 'num_iterations should equal to the length of epochs_per_iteration'
            self.epochs_per_iteration = epochs_per_iteration
        else:
            assert num_iterations > 0, 'num_iterations should >= 1'
            self.epochs_per_iteration = [epochs_per_iteration] * num_iterations

        self._validate_iteration_params()

        self._trainer = trainer
        self._criterion = criterion

    def _fresh_calculated(self):
        for wrapper in self.get_modules_wrapper():
            wrapper.if_calculated = False

    def _validate_iteration_params(self):
        assert all(num >= 0 for num in self.epochs_per_iteration), 'all epoch number need >= 0'

    def compress(self):
        training = self.bound_model.training
        self.bound_model.train()
        for _, epochs_num in enumerate(self.epochs_per_iteration):