Skip to content

PyTorch model conversion to ONNX, Keras, TFLite, CoreML

License

Notifications You must be signed in to change notification settings

opencv-ai/model_converter

Folders and files

NameName
Last commit message
Last commit date

Latest commit

Β 

History

30 Commits
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 

Repository files navigation

model_converter

Build status Python Version Dependencies Status

Code style: black Security: bandit Pre-commit Semantic Versions License Coverage Report

PyTorch model conversion toolbox

This toolbox supports model conversion to one of the following formats:

  • onnx
  • keras
  • tflite
  • coreml

Currently, two main conversion pipelines are supported:

  1. PyTorch --> ONNX --> Keras --> TFLite
  2. PyTorch --> TorchScript --> CoreML

Installation

Requirements

  • python 3.9

Install

It can be installed with the pip:

pip install git+ssh://[email protected]/opencv-ai/model_converter

Get started

To use converter in your project:

  1. Import converter:

    import model_converter
  2. Create an instance of a convertor:

    my_converter = model_converter.Converter(save_dir=<path to your output directory>, 
                                             simplify_exported_model=False
                                            )

    Use simplify_exported_model=True key to simplify onnx model.

  3. Run conversion of your model:

    converted_model = my_converter.convert(
        torch_model, # model for conversion
        torch_weights, # path to model checkpoint
        batch_size, # batch size
        input_size, # input size in [height, width] format
        channels, # number of input channels
        fmt, # output format for conversion - one of 'onnx', 'keras', 'tflite', 'coreml', 'tflite_coreml'
        force # set to `True` to rebuild all intermediate steps
    ) 

Model outputs wrapping (CoreML conversion)

You can wrap the output of your PyTorch model in a NamedTuple as shown below. By doing this, the Converter will be able to assign the correct names to the output in the resulting CoreML model.

class Model(nn.Module):
    """ 
    Parameters
    ----------
    nn : [nn.Module]
        Core feature extractor model that takes as input images and outputs feature 
        vector, e.g. of dimension Bx2048x7x7 
    """
    Output = collections.namedtuple('output', ['cls',])

    def __init__(self,
                 core: nn.Module):
        super().__init__()
        self.core = core

    def forward(self, x):
        return self.Output(cls=self.core(x))

Development

Poetry

Want to know more about Poetry? Check its documentation.

Details about Poetry

Poetry's commands are very intuitive and easy to learn, like:

  • poetry add numpy@latest
  • poetry run pytest
  • poetry publish --build

etc

Building and releasing your package

Building a new version of the application contains steps:

  • Bump the version of your package poetry version <version>. You can pass the new version explicitly, or a rule such as major, minor, or patch. For more details, refer to the Semantic Versions standard.
  • Make a commit to GitHub.
  • Create a GitHub release.
  • And... publish πŸ™‚ poetry publish --build

πŸš€ Features

Development features

Deployment features

Open source community features

Installation

pip install -U model_converter

or install with Poetry

poetry add model_converter

Makefile usage

Makefile contains a lot of functions for faster development.

1. Download and remove Poetry

To download and install Poetry run:

make poetry-download

To uninstall

make poetry-remove

2. Install all dependencies and pre-commit hooks

Install requirements:

make install

Pre-commit hooks coulb be installed after git init via

make pre-commit-install

3. Codestyle

Automatic formatting uses pyupgrade, isort and black.

make codestyle

# or use synonym
make formatting

Codestyle checks only, without rewriting files:

make check-codestyle

Note: check-codestyle uses isort, black and darglint library

Update all dev libraries to the latest version using one comand

make update-dev-deps
4. Code security

make check-safety

This command launches Poetry integrity checks as well as identifies security issues with Safety and Bandit.

make check-safety

5. Type checks

Run mypy static type checker

make mypy

6. Tests with coverage badges

Run pytest

make test

7. All linters

Of course there is a command to rule run all linters in one:

make lint

the same as:

make test && make check-codestyle && make mypy && make check-safety

8. Docker

make docker-build

which is equivalent to:

make docker-build VERSION=latest

Remove docker image with

make docker-remove

More information about docker.

9. Cleanup

Delete pycache files

make pycache-remove

Remove package build

make build-remove

Delete .DS_STORE files

make dsstore-remove

Remove .mypycache

make mypycache-remove

Or to remove all above run:

make cleanup

πŸ“ˆ Releases

You can see the list of available releases on the GitHub Releases page.

We follow Semantic Versions specification.

We use Release Drafter. As pull requests are merged, a draft release is kept up-to-date listing the changes, ready to publish when you’re ready. With the categories option, you can categorize pull requests in release notes using labels.

List of labels and corresponding titles

Label Title in Releases
enhancement, feature πŸš€ Features
bug, refactoring, bugfix, fix πŸ”§ Fixes & Refactoring
build, ci, testing πŸ“¦ Build System & CI/CD
breaking πŸ’₯ Breaking Changes
documentation πŸ“ Documentation
dependencies ⬆️ Dependencies updates

You can update it in release-drafter.yml.

GitHub creates the bug, enhancement, and documentation labels for you. Dependabot creates the dependencies label. Create the remaining labels on the Issues tab of your GitHub repository, when you need them.

πŸ›‘ License

License

This project is licensed under the terms of the MIT license. See LICENSE for more details.

πŸ“ƒ Citation

@misc{model_converter,
  author = {OpenCV.AI},
  title = {PyTorch model conversion to different formats},
  year = {2023},
  publisher = {GitHub},
  journal = {GitHub repository},
  howpublished = {\url{https://github.com/opencv-ai/model_converter}}
}

Credits πŸš€ Your next Python package needs a bleeding-edge project structure.

This project was generated with python-package-template

About

PyTorch model conversion to ONNX, Keras, TFLite, CoreML

Resources

License

Code of conduct

Security policy

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published