Open source PyTorch library to facilitate development and standardized evaluation of neural network pruning methods.
This repo contains the analysis and benchmarks results from the paper What is the State of Neural Network Pruning?.
First, install the dependencies, this repo depends on
PyTorch
Torchvision
NumPy
Pandas
Matplotlib
To install the dependencies
# Create a python virtualenv or conda env as necessary
# With conda
conda install numpy matplotlib pandas
conda install pytorch torchvision -c pytorch
# With pip
pip install numpy matplotlib pandas torch torchvision
then, to install the module itself you just need to clone the repo and add the parent path it to your PYTHONPATH
. For example:
git clone [email protected]:JJGO/shrinkbench.git shrinkbench
# Bash
echo "export PYTHONPATH=\"$PWD:\$PYTHONPATH\"" >> ~/.bashrc
# ZSH
echo "export PYTHONPATH=\"$PWD:\$PYTHONPATH\"" >> ~/.zshrc
ShrinkBench not only faciliates evaluation of pruning methods, but also their development. Here's the code for a simple implementation of Global Magnitude Pruning and Layerwise Magnitude Pruning. As you can see, it is quite succint; you are just tasked with implementing model_masks
a function that returns the masks for the model's weight tensors. If you want to prune your model layerwise, then you just need to implement layer_masks
. For more examples, see the source code for the provided baselines.
class GlobalMagWeight(VisionPruning):
def model_masks(self):
importances = map_importances(np.abs, self.params())
flat_importances = flatten_importances(importances)
threshold = fraction_threshold(flat_importances, self.fraction)
masks = importance_masks(importances, threshold)
return masks
class LayerMagWeight(LayerPruning, VisionPruning):
def layer_masks(self, module):
params = self.module_params(module)
importances = {param: np.abs(value) for param, value in params.items()}
masks = {param: fraction_mask(importances[param], self.fraction)
for param, value in params.items() if value is not None}
return masks
See here for a notebook showing how to run pruning experiments and plot their results
The modules are organized as follows:
submodule | Description |
---|---|
analysis/ |
Aggregated survey results over 80 pruning papers |
datasets/ |
Standardized dataloaders for supported datasets |
experiment/ |
Main experiment class with the data loading, pruning, finetuning & evaluation |
metrics/ |
Utils for measuring accuracy, model size, flops & memory footprint |
models/ |
Custom architectures not included in torchvision |
plot/ |
Utils for plotting across the logged dimensions |
pruning/ |
General pruning and masking API. |
scripts/ |
Executable scripts for running experiments (see experiment/ ) |
strategies/ |
Baselines pruning methods, mainly magnitude pruning based |