https://gist.github.com/SungFeng-Huang/57e1fce618d92b8f67350b31ed16ca5b
When using lottery ticket hypothesis, for each epoch, the unpruned parameters should be reset to its original values, while they actually didn't.
In my example, I print out the parameters after each forwarding. As it shows, the parameters at the start of the first training epoch is:
🐛 Bug
https://gist.github.com/SungFeng-Huang/57e1fce618d92b8f67350b31ed16ca5b When using lottery ticket hypothesis, for each epoch, the unpruned parameters should be reset to its original values, while they actually didn't. In my example, I print out the parameters after each forwarding. As it shows, the parameters at the start of the first training epoch is:
while for the second training epoch, the parameters are only pruned without reset:
The problem happens mainly due to L275: https://github.com/Lightning-AI/lightning/blob/07e7d6dc3ba1bbdae1a24da9d8d350096af68faa/src/pytorch_lightning/callbacks/pruning.py#L273-L279 assigning original values to
getattr(new, name)
is useless, sincegetattr(new, name)
would be further overwritten by the pruning function's forward_pre_hooks: https://github.com/pytorch/pytorch/blob/31142f57fc23edce291feaccf1670385e6239bbe/torch/nn/utils/prune.py#L23-L33To fix this problem, L275 should become:
To Reproduce
https://gist.github.com/SungFeng-Huang/57e1fce618d92b8f67350b31ed16ca5b
Expected behavior
The second (and later) epochs should start with parameters reset and pruned:
Environment
cc @carmocca