Simplification of pruned models for accelerated inference.
Fully supports torchvision classification models up to version 0.11.
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
For most scenarios the main simplify
function will suffice. This function returns the simplified model.
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. DefaultNone
.fuse_bn (bool)
: If True, fuse the conv-bn tuple.pinned_out (List)
: List ofnn.Modules
which output needs to remain of the original shape (e.g. layers related to a residual connection with a sum operation).
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)
The simplify
function is composed of three different submodules: fuse
, propagate
and remove
. Each module can be
used independently as needed.
Fuses adjacent Conv (or Linear) and BatchNorm layers.
Propagates non-zero bias of pruned neurons to remaining neurons of the next layers.
Removes zeroed neurons from the architecture.
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},
}