项目介绍
onnx2tflite 是一个供用户将 ONNX 格式的模型转换为 TensorFlow Lite 和 Keras 格式的工具。通过这些转换,用户可以方便地在不同的框架上使用同一个模型,从而实现模型的跨平台部署。项目欢迎大家参与讨论和提交修改建议。
安装指南
要开始使用 onnx2tflite,用户需要先从项目的 GitHub 仓库克隆代码,然后执行安装命令。具体步骤如下:
git clone https://github.com/MPolaris/onnx2tflite.git
cd onnx2tflite
python setup.py install
使用说明
用户可以通过以下命令将 ONNX 模型转换为 tflite 格式:
from onnx2tflite import onnx_converter
res = onnx_converter(
onnx_model_path = "./model.onnx",
need_simplify = True,
output_path = "./models/",
target_formats = ['tflite'],
)
为了更灵活的使用,onnx2tflite 提供了几种命令行格式来帮助用户指定不同的输出格式、保存路径以及量化配置等。
核心功能
- 高精度:转换后的模型与原始 ONNX 模型的平均误差小于每个元素的 (1e-5)。
- 转换效率:相较于其他工具,如 onnx_tf,该工具在转换为 TensorFlow Lite 模型时速度快了 30%。
- 自动通道对齐:自动将 PyTorch 格式 (NCWH) 转换为 TensorFlow 格式 (NWHC)。
- 量化支持:支持输出高精度 FP16 以及 UINT8 量化的模型。
- 代码简洁:保持代码结构简单明了,便于用户理解和使用。
从 Pytorch 到 TensorFlow Lite
用户可以直接使用 PyTorch 的 torchvision
模型,或者自定义模型,把它们转为 tflite 模型。以下是如何将 torchvision
中的模型转换为 tflite 的示例:
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
)
支持的模型
onnx2tflite 支持转换许多流行的深度学习模型,包括但不限于:
- SSD、HRNet、YOLOX、YOLOV3 到 YOLOV10 系列
- MoveNet、UNet/FPN、ViT、SwinTransformerV1
- MLP和DCGAN等自定义模型
- 所有
torchvision
分类模型及部分分割模型
自定义操作符
如果在使用过程中遇到不支持的操作符,用户可以自行添加支持,也可以在项目中提出问题。以上节选代码展示了如何通过名为 OPERATOR
的注册机制来添加一个名为 HardSigmoid
的操作符。
许可证
该软件遵循 Apache-2.0 许可证。用户可以在符合许可证条款的情况下自由使用和修改该软件。