VainF / Torch-Pruning

[CVPR 2023] Towards Any Structural Pruning; LLMs / SAM / Diffusion / Transformers / YOLOv8 / CNNs
https://arxiv.org/abs/2301.12900
MIT License
2.63k stars 329 forks source link

Feature: Improve Iterative Pruning: Verify Pruning Status Before Training #368

Open janthmueller opened 5 months ago

janthmueller commented 5 months ago

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.

janthmueller commented 3 months ago

This method is unreliable. Using the prune method doesn't consistently prune a group; sometimes nothing gets pruned. Therefore, an alternative solution is needed. Currently, I iterate over all parameters before and after pruning, which is inefficient but effective. We should consider another approach, possibly modifying the pruning history for stepwise pruning.