Simplify

Simplification of pruned models for accelerated inference.

Fully supports torchvision classification models up to version 0.11.

Installation

Simplify can be installed using pip:

pip3 install torch-simplify

or if you want to run the latest version of the code, you can install from git:

git clone https://github.com/EIDOSlab/simplify
cd simplify
pip3 install -r requirements.txt

Usage

Main function

For most scenarios the main simplify function will suffice. This function returns the simplified model.

Arguments

The expected arguments are:

model (torch.nn.Module): Module to be simplified i.e. the PyTorch's model.
x (torch.Tensor): zero-tensor of shape [1, C, N, M], same as the model usual input.
bn_folding (List): List of tuple (nn.Conv2d, nn.BatchNorm2d) to be fused. If None it tries to evaluate them given the model. Default None.
fuse_bn (bool): If True, fuse the conv-bn tuple.
pinned_out (List): List of nn.Modules which output needs to remain of the original shape (e.g. layers related to a residual connection with a sum operation).

Minimal working example

import torch
from torchvision import models
from simplify import simplify

model = models.resnet18()

# Apply some pruning strategy or load a pruned checkpoint

dummy_input = torch.zeros(1, 3, 224, 224)  # Tensor shape is that of a standard input for the given model
simplified_model = simplify(model, dummy_input)

Submodules

The simplify function is composed of three different submodules: fuse, propagate and remove. Each module can be used independently as needed.

fuse

Fuses adjacent Conv (or Linear) and BatchNorm layers.

propagate

Propagates non-zero bias of pruned neurons to remaining neurons of the next layers.

remove

Removes zeroed neurons from the architecture.

Citing

If you use this software for research or application purposes, please use the following citation:

@article{bragagnolo2021simplify,
  title = {Simplify: A Python library for optimizing pruned neural networks},
  journal = {SoftwareX},
  volume = {17},
  pages = {100907},
  year = {2022},
  issn = {2352-7110},
  doi = {https://doi.org/10.1016/j.softx.2021.100907},
  url = {https://www.sciencedirect.com/science/article/pii/S2352711021001576},
  author = {Andrea Bragagnolo and Carlo Alberto Barbano},
}

EIDOSLAB/simplify