ONNX转Keras和ONNX转TFLite工具

欢迎

如果您有任何好的想法，欢迎讨论或提交项目PR。

安装

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

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

---

# 基础用法
python -m onnx2tflite --weights "./your_model.onnx"

# 指定保存路径
python -m onnx2tflite --weights "./your_model.onnx" --outpath "./save_path"

# 仅保存TFLite模型
python -m onnx2tflite --weights "./your_model.onnx" --outpath "./save_path" --formats "tflite"

# 同时保存Keras和TFLite模型
python -m onnx2tflite --weights "./your_model.onnx" --outpath "./save_path" --formats "tflite" "keras"

# 截断模型、重新定义输入输出，并支持中间层
python -m onnx2tflite --weights "./your_model.onnx" --outpath "./save_path" --formats "tflite" --input-node-names "layer_inputname" --output-node-names "layer_outname1" "layer_outname2"

# 对模型权重进行量化，仅量化权重
python -m onnx2tflite --weights "./your_model.onnx" --formats "tflite" --weigthquant

# 对模型权重进行量化，同时包括输入和输出
## FP16量化
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
## 生成随机数据，而非从图像文件中读取
python -m onnx2tflite --weights "./your_model.onnx" --formats "tflite" --int8

---

特性

• 高一致性。与ONNX输出相比，每个元素的平均误差小于1e-5。
• 更快的速度。生成TensorFlow Lite模型的速度比onnx_tf快30%。
• 自动通道对齐。自动将PyTorch格式(NCWH)转换为TensorFlow格式(NWHC)。
• 部署支持。支持输出量化模型，包括FP16量化和UINT8量化。
• 代码友好。我一直致力于保持代码结构简单明了。

---

PyTorch -> ONNX -> TensorFlow-Keras -> TensorFlow-Lite

• 从torchvision到tensorflow-lite

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)# 或 opset_version=13

from converter import onnx_converter
onnx_converter(
    onnx_model_path = "./mobilenetV2.onnx",
    need_simplify = True,
    output_path = "./",
    target_formats = ['tflite'], # 或 ['keras'], ['keras', 'tflite']
    weight_quant = False,
    fp16_model=False,
    int8_model = False,
    int8_mean = None,
    int8_std = None,
    image_root = None
)

• 从自定义PyTorch模型到tensorflow-lite-int8

import torch
import torch.nn as nn
import torch.nn.functional as F

class MyModel(nn.Module):
    def __init__(self):
        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)# 或 opset_version=13

from converter import onnx_converter
onnx_converter(
    onnx_model_path = "./mymodel.onnx",
    need_simplify = True,
    output_path = "./",
    target_formats = ['tflite'], #或 ['keras'], ['keras', '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" # 给出训练集的图像文件夹
)

---

已验证的模型

• SSD
• HRNet
• YOLOX
• YOLOV3
• YOLOV4
• YOLOV5
• YOLOV6
• YOLOV7
• YOLOV10
• MoveNet
• UNet\FPN
• ViT(torchvision)
• SwinTransformerV1
• MLP(自定义)
• DCGAN(自定义)
• AutoEncoder/VAE
• 所有torchvision分类模型
• torchvision中的一些分割模型
• 无特殊算子的1D或2D CNN(自定义)

---

自行添加算子

当遇到不支持的算子时，您可以选择自行添加该算子，或提交问题。
按照以下步骤实现新的算子解析器非常简单。
步骤0：在layers文件夹中选择对应的层代码文件，例如‘HardSigmoid’对应activations_layers.py。
步骤1：打开并编辑该文件：

# 所有算子都通过OPERATOR注册。
# 注册的算子名称即为ONNX算子名称。
@OPERATOR.register_operator("HardSigmoid")
class TFHardSigmoid():
    def __init__(self, tensor_grap, node_weights, node_inputs, node_attribute, node_outputs, layout_dict, *args, **kwargs) -> None:
        '''
        :param tensor_grap: 字典，键为节点名称，值为TensorFlow-Keras节点的输出张量。
        :param node_weights: 字典，键为节点名称，值为静态数据，如权重/偏置/常数；权重通常需要通过dimension_utils.tensor_NCD_to_NDC_format函数进行转换。
        :param node_inputs: 字符串列表，记录输入节点的名称，表明输入来自哪些节点，可能是tensor_grap或node_weights中的内容。
        :param node_attribute: 字典，键为属性名称，如'axis'或'perm'，值类型不确定，可能是List[int]、int或float。注意，'axis'的值应通过dimension_utils.channel_to_last_dimension或dimension_utils.shape_NCD_to_NDC_format函数从NCHW格式调整为NHWC格式。
        :param node_inputs: 字符串列表，记录输出节点的名称。
        :param layout_dict: 列表，记录所有前置节点的布局。
        '''
        super().__init__()
        self.alpha = node_attribute.get("alpha", 0.2)
        self.beta = node_attribute.get("beta", 0.5)

    def __call__(self, inputs):
        return tf.clip_by_value(self.alpha*inputs+self.beta, 0, 1)

步骤2：确保代码运行无误。
步骤3：将模型转换为TFLite，无需任何量化。

---

许可证

本软件受Apache-2.0许可证保护。

onnx2tflite 快速上手指南

onnx2tflite 是一款高效工具，用于将 ONNX 模型转换为 TensorFlow Keras 或 TensorFlow Lite (TFLite) 格式。它支持自动通道对齐（NCHW 转 NHWC）、高精度转换（误差 < 1e-5）以及多种量化方案（FP16/INT8），特别适合将 PyTorch 模型部署到移动端。

环境准备

• 操作系统: Linux, macOS, Windows
• Python 版本: 建议 Python 3.6+
• 前置依赖:
  • tensorflow (建议 2.x 版本)
  • onnx
  • onnx-simplifier (可选，用于模型简化)
  • torch / torchvision (如果需要从 PyTorch 导出 ONNX)

提示：国内开发者安装依赖时，推荐使用清华或阿里镜像源加速：

pip install tensorflow onnx onnx-simplifier -i https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

通过源码克隆并安装：

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

基本使用

1. Python API 调用（推荐）

这是最简洁的使用方式，适合集成到自动化脚本中。以下示例将 model.onnx 转换为 TFLite 格式，并自动进行模型简化。

from onnx2tflite import onnx_converter

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

2. 命令行工具 (CLI)

如果你更喜欢命令行操作，可以使用以下基础命令：

基础转换：

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

指定输出路径和格式：

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

3. 典型工作流：PyTorch -> ONNX -> TFLite

对于 PyTorch 用户，完整的转换流程如下：

1. 导出 ONNX:

   import torch
   import torchvision
   
   # 加载模型
   model = torchvision.models.mobilenet_v2(True)
   _input = torch.randn(1, 3, 224, 224)
   
   # 导出为 ONNX (opset_version 建议 11 或 13)
   torch.onnx.export(model, _input, './mobilenetV2.onnx', opset_version=11)
   

2. 转换为 TFLite:

   from onnx2tflite import onnx_converter
   
   onnx_converter(
       onnx_model_path = "./mobilenetV2.onnx",
       need_simplify = True,
       output_path = "./",
       target_formats = ['tflite'],
       weight_quant = False,   # 是否仅权重量化
       fp16_model=False,       # 是否 FP16 量化
       int8_model = False,     # 是否 INT8 量化
   )
   

进阶提示：若需进行 INT8 量化以提升推理速度，需在 onnx_converter 中设置 int8_model=True，并提供训练集路径 (image_root) 及预处理均值/标准差 (int8_mean, int8_std)。