Feature: Improve Iterative Pruning: Verify Pruning Status Before Training

[janthmueller]

Hello, I've noticed that during iterative pruning, the model might not necessarily undergo pruning after each step (pruning generator does not return anything - is empty). Since the iterative pruning process involves a cycle of pruning and training, it would be beneficial to verify whether the model has been pruned after each step, rather than proceeding directly to training.

Here's a scenario where this improvement could be useful: currently, after calling the step() function of the pruning algorithm, we proceed to training the model without confirming if pruning has actually occurred. This can lead to unnecessary training cycles on an unpruned model.

To address this, we can modify the step() method within the MetaPruner class. By introducing a return value to indicate whether pruning has taken place, we can optimize the training process. Below is a simple suggested implementation:

def step(self, interactive=False) -> typing.Union[typing.Generator, None]:
    self.current_step += 1
    pruning_method = self.prune_global if self.global_pruning else self.prune_local

    if interactive: 
        return pruning_method() # yield groups for interactive pruning
    else:
        pruned = False
        for group in pruning_method():
            group.prune()
            pruned = True
        return pruned

With this enhancement, before initiating training, we can easily check if pruning has occurred after calling step(). This allows us to seamlessly continue our iterative loop without unnecessary training cycles on an unpruned model.