onnx2tflite

ONNX2TFLite Project Overview

Welcome

The ONNX2TFLite project invites discussions and contributions from the community. Developers with innovative ideas are encouraged to contribute to the project's growth and success through pull requests.

Installation

To get started with ONNX2TFLite, users can easily clone the repository and install the necessary packages by following these commands:

git clone https://github.com/MPolaris/onnx2tflite.git
cd onnx2tflite
python setup.py install

Once installed, converting an ONNX model to TF-Lite can be done with a few lines of code:

from onnx2tflite import onnx_converter
res = onnx_converter(
        onnx_model_path = "./model.onnx",
        need_simplify = True,
        output_path = "./models/",
        target_formats = ['tflite'],
    )

Command Line Usage

Users can also convert models using the command line with various options tailored to their needs. Below are some examples of how to execute these commands:

• Basic Conversion

  python -m onnx2tflite --weights "./your_model.onnx"
  

• Specify Save Path

  python -m onnx2tflite --weights "./your_model.onnx" --outpath "./save_path"
  

• Save as TFLite

  python -m onnx2tflite --weights "./your_model.onnx" --outpath "./save_path" --formats "tflite"
  

• Save both Keras and TFLite Models

  python -m onnx2tflite --weights "./your_model.onnx" --outpath "./save_path" --formats "tflite" "keras"
  

• Redefine Model Inputs and Outputs

  python -m onnx2tflite --weights "./your_model.onnx" --outpath "./save_path" --formats "tflite" --input-node-names "layer_inputname" --output-node-names "layer_outname1" "layer_outname2"
  

• Quantify Model Weights

  • Only Weights

    python -m onnx2tflite --weights "./your_model.onnx" --formats "tflite" --weigthquant
    

  • Including Inputs and Outputs

    python -m onnx2tflite --weights "./your_model.onnx" --formats "tflite" --fp16
    python -m onnx2tflite --weights "./your_model.onnx" --formats "tflite" --int8 --imgroot "./dataset_path" --int8mean 0 0 0 --int8std 255 255 255
    

Features

• High Consistency: Ensures average error is less than 1e-5 per element compared to ONNX outputs.
• Faster Processing: Produces TensorFlow Lite models 30% faster than onnx_tf.
• Automatic Channel Alignment: Converts PyTorch NCWH format to TensorFlow NWHC format automatically.
• Deployment Support: Supports outputting quantified models, including options for fp16 and uint8 quantization.
• Developer-Friendly: The code is structured to be straightforward and easy to understand.

Model Conversion Process

Pytorch to TensorFlow Lite

Users can convert models from PyTorch to TensorFlow Lite, supporting int8 quantization, through straightforward Python scripts.

• Example with torchvision models:

  import torch
  import torchvision
  
  _input = torch.randn(1, 3, 224, 224)
  model = torchvision.models.mobilenet_v2(True)
  
  torch.onnx.export(model, _input, './mobilenetV2.onnx', opset_version=11)
  
  from converter import onnx_converter
  onnx_converter(
      onnx_model_path = "./mobilenetV2.onnx",
      need_simplify = True,
      output_path = "./",
      target_formats = ['tflite'],
      weight_quant = False,
      fp16_model=False,
      int8_model = False
  )
  

• Example with a custom PyTorch model:

  import torch
  import torch.nn as nn
  
  class MyModel(nn.Module):
      def __init__(self):
          super(MyModel, self).__init__()
          self.conv = nn.Sequential(
              nn.Conv2d(3, 64, kernel_size=3, padding=1),
              nn.BatchNorm2d(64),
              nn.ReLU(inplace=True),
          )
      
      def forward(self, x):
          return self.conv(x)
  
  model = MyModel()
  model.load_state_dict(torch.load("model_checkpoint.pth", map_location="cpu"))
  
  _input = torch.randn(1, 3, 224, 224)
  torch.onnx.export(model, _input, './mymodel.onnx', opset_version=11)
  
  from converter import onnx_converter
  onnx_converter(
      onnx_model_path = "./mymodel.onnx",
      need_simplify = True,
      output_path = "./",
      target_formats = ['tflite'],
      weight_quant = False,
      int8_model = True,
      int8_mean = [123.675, 116.28, 103.53],
      int8_std = [58.395, 57.12, 57.375],
      image_root = "./dataset/train"
  )
  

Validated Models

The tool has been tested with numerous models, including SSD, HRNet, various YOLO versions (YOLOX, YOLOV3, etc.), MoveNet, UNet/FPN, SwinTransformerV1, and many more from torchvision.

Adding Custom Operators

For unsupported operations, developers can easily add their own implementations. The process involves registering the operator with a simple workflow, making modifications to existing layer files, and testing the conversion to ensure it functions correctly.

License

The ONNX2TFLite tool is licensed under the Apache-2.0 license, which allows users wide freedom to use, modify, and distribute the tool and any derivative works.