tensorflow - 你如何将 .onnx 转换为 tflite？

Question

我已通过以下方式将模型导出到 ONNX:

# Export the model
torch_out = torch.onnx._export(learn.model,             # model being run
                           x,                       # model input (or a tuple for multiple inputs)
                          EXPORT_PATH + "mnist.onnx", # where to save the model (can be a file or file-like object)
                           export_params=True)      # store the trained parameter weights inside the model file

现在我正在尝试将模型转换为 Tensorflow Lite 文件，以便我可以在 Android 上进行推理。不幸的是，PyTorch/Caffe2 对 Android 的支持相当缺乏或过于复杂，但 Tensorflow 看起来要简单得多。

ONNX 到 Tflite 的文档对此非常清楚。

我尝试通过以下方式导出到 Tensorflow GraphDef proto:
tf_rep.export_graph(EXPORT_PATH + 'mnist-test/mnist-tf-export.pb')
然后运行 ​​ toco :

toco \
--graph_def_file=mnist-tf-export.pb \
--input_format=TENSORFLOW_GRAPHDEF \
--output_format=TFLITE \
--inference_type=FLOAT \
--input_type=FLOAT \
--input_arrays=0 \
--output_arrays=add_10 \
--input_shapes=1,3,28,28 \
--output_file=mnist.tflite`

当我这样做时，虽然我收到以下错误:

File "anaconda3/lib/python3.6/site-packages/tensorflow/lite/python/convert.py", line 172, in toco_convert_protos
    "TOCO failed. See console for info.\n%s\n%s\n" % (stdout, stderr))
tensorflow.lite.python.convert.ConverterError: TOCO failed. See console for info.
2018-11-06 16:28:33.864889: I tensorflow/lite/toco/import_tensorflow.cc:1268] Converting unsupported operation: PyFunc
2018-11-06 16:28:33.874130: F tensorflow/lite/toco/import_tensorflow.cc:114] Check failed: attr.value_case() == AttrValue::kType (1 vs. 6)

此外，即使我运行命令时，我也不知道要为 input_arrays 或 output_arrays 指定什么，因为模型最初是在 PyTorch 中构建的。

有没有人成功地将他们的 ONNX 模型转换为 TFlite？

这是我要转换的 ONNX 文件: https://drive.google.com/file/d/1sM4RpeBVqPNw1WeCROpKLdzbSJPWSK79/view?usp=sharing

额外信息

Python 3.6.6::Anaconda 自定义(64 位)

昂克斯。 版本 = '1.3.0'

tf。 版本 = '1.13.0-dev20181106'

火炬。 版本 = '1.0.0.dev20181029'

Answer 1

我认为 ONNX 文件，即 model.onnx您提供的已损坏我不知道是什么问题，但它没有对 ONNX 运行时进行任何推断。

Now you can run PyTorch Models directly on mobile phones. check out PyTorch Mobile's documentation here

This answer is for TensorFlow version 1,
For TensorFlow version 2 or higher click link

将模型从 protobuf freezeGraph 转换为 TFlite 的最佳方法是使用官方的 TensorFlow lite 转换器 documentation
根据 TensorFlow Docs , TocoConverter 已被弃用

This class (tf.compat.v1.lite.TocoConverter) has been deprecated. Please use lite.TFLiteConverter instead.

从 PyTorch 转换为 ONNX 模型
从 转换模型的最佳实践pytorch 至 Onnx 就是你应该在 中添加以下参数来指定你的模型的输入和输出层的名称torch.onnx.export() 功能

# Export the model from PyTorch to ONNX
torch_out = torch.onnx._export(model,             # model being run
                                x,          # model input (or a tuple for multiple inputs)
                                EXPORT_PATH + "mnist.onnx",      # where to save the model (can be a file or file-like object)
                                export_params=True,       # store the trained parameter weights inside the model file
                                input_names=['main_input'],     # specify the name of input layer in onnx model
                                output_names=['main_output'])     # specify the name of input layer in onnx model

所以在你的情况下:
现在使用 onnx-tf 将此模型导出到 TensorFlow protobuf FreezeGraph

Please note that this method is only working when tensorflow_version < 2

从 ONNX 转换为 TensorFlow freezGraph
要转换模型，请从以下命令安装 onnx-tf 版本 1.5.0

pip install  onnx-tf==1.5.0

现在要将 .onnx 模型转换为 TensorFlow 卡住图，请在 shell 中运行以下命令

onnx-tf convert -i "mnist.onnx" -o  "mnist.pb"

从 TensorFlow FreezeGraph .pb 转换为 TF
现在将此模型从 .pb 文件转换为 tflite 模型使用此代码

import tensorflow as tf
# make a converter object from the saved tensorflow file
converter = tf.lite.TFLiteConverter.from_frozen_graph('mnist.pb', #TensorFlow freezegraph .pb model file
                                                      input_arrays=['main_input'], # name of input arrays as defined in torch.onnx.export function before.
                                                      output_arrays=['main_output']  # name of output arrays defined in torch.onnx.export function before.
                                                      )
# tell converter which type of optimization techniques to use
converter.optimizations = [tf.lite.Optimize.DEFAULT]
# to view the best option for optimization read documentation of tflite about optimization
# go to this link https://www.tensorflow.org/lite/guide/get_started#4_optimize_your_model_optional

# convert the model 
tf_lite_model = converter.convert()
# save the converted model 
open('mnist.tflite', 'wb').write(tf_lite_model)

要为您的模型用例选择最适合优化的选项，请参阅有关 TensorFlow lite 优化的官方指南
https://www.tensorflow.org/lite/guide/get_started#4_optimize_your_model_optional

Note: You can try my Jupyter Notebook Convert ONNX model to Tensorflow Lite on Google Colaboratory link