onnx2tflite: 高效转换ONNX模型到TensorFlow Lite的开源工具

Ray

onnx2tflite

onnx2tflite:连接ONNX与TensorFlow Lite的桥梁

在当今深度学习蓬勃发展的时代,不同深度学习框架之间的模型转换一直是一个重要而富有挑战性的问题。ONNX(开放神经网络交换)作为一种通用的深度学习模型表示格式,为不同框架之间的模型交换提供了可能。而TensorFlow Lite作为一个轻量级的移动端和嵌入式设备推理框架,在模型部署方面发挥着重要作用。onnx2tflite项目就是为了解决ONNX模型到TensorFlow Lite模型的转换而生的。

onnx2tflite的主要特性

onnx2tflite是一个功能强大、使用便捷的开源工具,它具有以下几个主要特性:

高度一致性: 与ONNX模型的输出相比,转换后的模型平均每个元素的误差小于1e-5,保证了转换的精度。
更快的转换速度: 与onnx_tf等其他工具相比,onnx2tflite在输出TensorFlow Lite模型时速度提升了30%。
自动通道对齐: 能够自动将PyTorch格式(NCWH)转换为TensorFlow格式(NWHC),解决了不同框架间的数据格式差异问题。
部署支持: 支持输出量化模型,包括fp16量化和uint8量化,为模型在资源受限设备上的部署提供了便利。
代码友好: 项目的代码结构简单清晰,便于理解和扩展。

安装与基本使用

要使用onnx2tflite,首先需要从GitHub克隆项目并安装:

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

安装完成后,可以通过以下方式使用onnx2tflite:

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

这段代码将一个ONNX模型转换为TensorFlow Lite模型。need_simplify参数表示是否需要对模型进行简化,target_formats参数指定了输出格式。

命令行使用方法

除了作为Python模块使用,onnx2tflite还提供了命令行接口,使用更加灵活:

# 基本使用
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"

高级功能:模型剪枝与量化

onnx2tflite还支持一些高级功能,如模型剪枝和量化:

模型剪枝: 可以重新定义模型的输入输出,支持中间层:

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

模型量化: 支持权重量化和完整的int8量化:

# 仅权重量化
python -m onnx2tflite --weights "./your_model.onnx" --formats "tflite" --weigthquant

# fp16量化
python -m onnx2tflite --weights "./your_model.onnx" --formats "tflite" --fp16

# int8量化(推荐)
python -m onnx2tflite --weights "./your_model.onnx" --formats "tflite" --int8 --imgroot "./dataset_path" --int8mean 0 0 0 --int8std 255 255 255

支持的模型类型

onnx2tflite已经在多种常见的深度学习模型上进行了验证,包括但不限于:

目标检测模型:SSD, YOLOX, YOLOV3/4/5/6/7/10
姿态估计模型:MoveNet
图像分割模型:UNet, FPN
图像分类模型:ViT, SwinTransformerV1
生成模型:DCGAN, AutoEncoder/VAE
自定义的1D或2D CNN模型

此外,torchvision中的所有分类模型和部分分割模型也都得到了支持。这种广泛的模型支持使得onnx2tflite成为一个非常实用的工具,能够满足大多数深度学习从业者的需求。

自定义算子支持

对于onnx2tflite尚未支持的算子,用户可以自行添加。项目提供了简单的步骤来实现新的算子解析器:

在layers文件夹中选择合适的文件(如activations_layers.py)。
使用@OPERATOR.register_operator装饰器注册新的算子。
实现算子的__init__和__call__方法。

例如,添加HardSigmoid算子的示例代码如下:

@OPERATOR.register_operator("HardSigmoid")
class TFHardSigmoid():
    def __init__(self, tensor_grap, node_weights, node_inputs, node_attribute, node_outputs, layout_dict, *args, **kwargs) -> None:
        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)

这种灵活的扩展机制使得onnx2tflite能够不断适应新的模型结构和算子。

onnx2tflite的优势

简化工作流程: onnx2tflite提供了一站式的解决方案,从ONNX直接转换到TensorFlow Lite,简化了模型转换的工作流程。
性能优化: 通过支持多种量化方式,onnx2tflite可以帮助用户获得更小、更快的模型,适合在移动设备和嵌入式系统上部署。
广泛兼容性: 支持多种常见的深度学习模型,覆盖了计算机视觉、自然语言处理等多个领域。
易于使用: 提供了直观的Python API和命令行接口,使用门槛低,适合不同背景的开发者。
开源社区支持: 作为一个开源项目,onnx2tflite得到了社区的支持和贡献,不断改进和更新。

结语

onnx2tflite为深度学习模型的转换和部署提供了一个强大而灵活的工具。它不仅简化了从ONNX到TensorFlow Lite的转换过程,还提供了多种优化选项,帮助开发者更好地将模型部署到资源受限的设备上。随着深度学习技术的不断发展和普及,onnx2tflite这样的工具将在模型部署和优化中发挥越来越重要的作用。

无论您是研究人员、开发者还是AI从业者,onnx2tflite都是一个值得尝试的工具。它可以帮助您更有效地管理和部署深度学习模型,为您的AI项目增添新的可能性。我们期待看到更多基于onnx2tflite的创新应用,也欢迎社区成员为这个开源项目做出贡献,共同推动深度学习技术的发展与应用。

🔗 项目链接: onnx2tflite GitHub仓库