Skip to content

Talmaj/onnx2pytorch

Repository files navigation

ONNX to PyTorch

PyPI - License CircleCI Downloads PyPI

A library to transform ONNX model to PyTorch. This library enables use of PyTorch backend and all of its great features for manipulation of neural networks.

Installation

pip install onnx2pytorch

Usage

import onnx
from onnx2pytorch import ConvertModel

onnx_model = onnx.load(path_to_onnx_model)
pytorch_model = ConvertModel(onnx_model)

Currently supported and tested models from onnx_zoo:

Limitations

Known current version limitations are:

  • batch_size > 1 could deliver unexpected results due to ambiguity of onnx's BatchNorm layer.
    That is why in this case for now we raise an assertion error.
    Set experimental=True in ConvertModel to be able to use batch_size > 1.
  • Fine tuning and training of converted models was not tested yet, only inference.

Development

Dependency installation

pip install -r requirements.txt

From onnxruntime>=1.5.0 you need to add the following to your .bashrc or .zshrc if you are running OSx: export KMP_DUPLICATE_LIB_OK=True

Code formatting

The Uncompromising Code Formatter: Black
black {source_file_or_directory}

Install it into pre-commit hook to always commit nicely formatted code:
pre-commmit install

Testing

Pytest and tox.
tox

Test fixtures

To test the complete conversion of an onnx model download pre-trained models: ./download_fixtures.sh
Use flag --all to download more models. Add any custom models to ./fixtures folder to test their conversion.

Debugging

Set ConvertModel(..., debug=True) to compare each converted activation from pytorch with the activation from onnxruntime.
This helps identify where in the graph the activations start to differ.