Skip to main content

Simplification of pruned models for accelerated inference

Project description


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
cd simplify
pip3 install -r requirements.txt


Main function

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. 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)


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:

  title = {Simplify: A Python library for optimizing pruned neural networks},
  journal = {SoftwareX},
  volume = {17},
  pages = {100907},
  year = {2022},
  issn = {2352-7110},
  doi = {},
  url = {},
  author = {Andrea Bragagnolo and Carlo Alberto Barbano},

Project details

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

torch-simplify-1.1.4.tar.gz (12.9 kB view hashes)

Uploaded source

Built Distribution

torch_simplify-1.1.4-py3-none-any.whl (16.5 kB view hashes)

Uploaded py3

Supported by

AWS AWS Cloud computing and Security Sponsor Datadog Datadog Monitoring Fastly Fastly CDN Google Google Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Sentry Sentry Error logging StatusPage StatusPage Status page