mil-ad / prospr

Prospect Pruning: Finding Trainable Weights at Initialization Using Meta-Gradients

Prospect Pruning (ProsPr)

The code for "Prospect Pruning: Finding Trainable Weights at Initialization Using Meta-Gradients"

Installation

1️⃣ Reproducing results

You can replicate the development environment and use the same models and training script used in the paper with:

$ conda env create -f environment.yml

If you'd like to use the exact same package versions we used:

$ conda env create -f environment_pinned.yml

This will create the Conda environment prospr. The project's entry point is cli.py

To see the available options and switches:

$ python cli.py -h

2️⃣ As a package

You can also install and use ProsPr as a package inside your own projects:

$ pip install git+ssh://[e-mail]/mil-ad/prospr.git

The prospr package can then be imported and used:

import prospr

help(prospr)

pruned_model = prospr.prune(
    model,
    prune_ratio=0.98,
    dataloader=train_dataloader,
    filter_fn=prune_filter_fn,
    num_steps=3,
    inner_lr=0.5,
    inner_momentum=0.9,
)

Citation

@article{alizadeh2022prospect,
  title = {Prospect Pruning: Finding Trainable Weights at Initialization using Meta-Gradients},
  author = {Alizadeh, Milad and Tailor, Shyam A. and Zintgraf, Luisa M and van Amersfoort, Joost and Farquhar, Sebastian and Lane, Nicholas Donald and Gal, Yarin},
  booktitle = {International Conference on Learning Representations},
  year = {2022}
}