Skip to content
Simplify your onnx model
Python
Branch: master
Clone or download
Fetching latest commit…
Cannot retrieve the latest commit at this time.
Permalink
Type Name Latest commit message Commit time
Failed to load latest commit information.
.github/ISSUE_TEMPLATE Update issue templates Sep 18, 2019
imgs Move images in README.md to imgs/ Apr 2, 2019
onnxsim Add some comments for last commit Aug 6, 2019
.gitignore Update .gitignore Apr 2, 2019
.travis.yml Set version in setup.py in CI Jul 29, 2019
LICENSE Create LICENSE Apr 2, 2019
README.md
setup.py Set version in setup.py in CI Jul 29, 2019

README.md

ONNX Simplifier

PyPI version PyPI pyversions PyPI license PRs Welcome

ONNX is great, but sometimes too complicated.

Background

One day I wanted to export the following simple reshape operation to ONNX:

import torch


class JustReshape(torch.nn.Module):
    def __init__(self):
        super(JustReshape, self).__init__()

    def forward(self, x):
        return x.view((x.shape[0], x.shape[1], x.shape[3], x.shape[2]))


net = JustReshape()
model_name = 'just_reshape.onnx'
dummy_input = torch.randn(2, 3, 4, 5)
torch.onnx.export(net, dummy_input, model_name, input_names=['input'], output_names=['output'])

Input shape in ONNX is static, so what I expected is

simple_reshape

However, I got the following complicated model even after polishing:

complicated_reshape

Moreover, there are also some operations performed on weights (like this), which can all be eliminated by offline computation.

Our solution

ONNX Simplifier is presented to simplify the ONNX model. It infers the whole computation graph and then replaces the redundant operators with their constant outputs.

Install it via pip (Python >= 3.5)

pip3 install onnx-simplifier

Then

python3 -m onnxsim input_model output_model

Demonstration

An overall comparison between a complicated model and its simplified version:

Comparison between old model and new model

You can’t perform that action at this time.