ShrinkBench

Open source PyTorch library to facilitate development and standardized evaluation of neural network pruning methods.

Paper

This repo contains the analysis and benchmarks results from the paper What is the State of Neural Network Pruning?.

Installation

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 tqdm
conda install pytorch torchvision -c pytorch

# With pip
pip install numpy matplotlib pandas torch torchvision tqdm

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 git@github.com:JJGO/shrinkbench.git shrinkbench

# Bash
echo "export PYTHONPATH=\"$PWD:\$PYTHONPATH\"" >> ~/.bashrc

# ZSH
echo "export PYTHONPATH=\"$PWD:\$PYTHONPATH\"" >> ~/.zshrc

Strategies

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

Experiments

See here for a notebook showing how to run pruning experiments and plot their results

Modules

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