search

TimDettmers / sparse_learning

Sparse learning library and sparse momentum resources.
MIT License
377 stars 45 forks source link

Masking.prune_every_k_steps default value #8

Closed evcu closed 5 years ago

evcu commented 5 years ago

Hi Tim,

I was looking at your code and couldn't find where the prune_every_k_steps is set. It is None by default, which prevents pruning/growing.

Thanks

TimDettmers commented 5 years ago

This is a hidden feature I played around with on ImageNet to prune more than once per epoch. It is fully functional but it is currently not used in the main code. Currently, pruning/growing is done via mask.at_end_of_epoch() which handles all the details and mask.step() does all the masking of the weights.

However, if you want to use prune_every_k_steps you can just initialize it to a value, for example mask.prune_every_k_steps = 540 would prune every 540 mini-batches which for MNIST with batch size 100 would be exactly one epoch. If you want to prune twice per epoch, just set it to 270 etcetera. This is done automatically and you just need to use mask.step(). If you do this make sure to not use it together with mask.at_end_of_epoch().

evcu commented 5 years ago

Right! Now see it. Thanks for the rapid response.