Pytorch转tensorflow、tflite，ONNX转为tensorflow、tflite模型

当你训练了一个Pytorch的模型或者在某处找到了一个ONNX模型，又需要将其转化为Keras模型或者tensorflow lite进行部署时，可以使用两种方式：

1. ONNX提供的onnx_tf(https://github.com/onnx/onnx-tensorflow)

该种方式具有的优点的是支持的算法更为的丰富，转换成功率更高，不过缺点会出现很多无用的transpose操作导致模型运行速度较慢。

2. github社区的开源项目onnx2tflite(https://github.com/MPolaris/onnx2tflite)

该项目具有以下优点

• 高一致性。相比于ONNX模型的输出，每一位的平均误差小于1e-5。
• 更快的运行速度。相比于onnx_tf输出的模型，运行速度快30%。
• 自动通道变换。能够将NCWH格式的通道变换为tensorflow的NWHC格式。
• 支持模型量化。继承了tensorflow lite量化工具，支持模型权重量化以及整体uint8量化，模型更小、运行速度更快。
• 代码更友好。个人认为我的代码逻辑较为清晰，更方便开发者进行更改。

目前该项目以及支持并验证了以下模型

• SSD
• UNet
• HRNet
• MoveNet
• YOLO 系列，V3\V4\V5\V6\V7，以及YOLOX均已支持
• torchvision种所有的分类模型、以及部分分割模型
• 一些常规的2D卷积神经网络

3. 一个Pytorch模型的转换小例子

# git clone https://github.com/MPolaris/onnx2tflite
# cd onnx2tflite
# pip install -r requirements.txt

import torch
import torchvision
_input = torch.randn(1, 3, 224, 224)
model = torchvision.models.mobilenet_v2(True)
# use default settings is ok
torch.onnx.export(model, _input, './mobilenetV2.onnx', opset_version=11)# or opset_version=13

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

结语

目前onnx2tflite工具还在持续更新，希望大家能够给这个项目一个star，在使用过程中遇到什么问题欢迎提issue和pr。