您如何将.onnx转换为tflite?

时间:2018-11-07 00:42:40

标签: tensorflow keras deep-learning pytorch onnx

我已通过以下方式将模型导出到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上进行推理。不幸的是,对于Android来说,PyTorch / Caffe2的支持还相当缺乏或过于复杂,但是Tensorflow看起来要简单得多。

ONNX to Tflite的文档对此很浅。

我尝试通过以下方式导出到Tensorflow GraphDef原型:

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位)
  • onnx。版本 ='1.3.0'
  • tf。版本 ='1.13.0-dev20181106'
  • 火炬。版本 ='1.0.0.dev20181029'

2 个答案:

答案 0 :(得分:1)

我认为您提供的onnx文件已损坏,我不知道这是什么问题,但它没有对onnx运行时进行任何推断
将模型从 Pytorch 转换为 Onnx 的最佳做法是,应添加以下参数以在 torch.onnx中指定模型的输入和输出层的名称。 export()函数
因此,在您的情况下:

var playWins = 0;
var compWins = 0;

现在使用onnx-tf将此模型导出到tensorFlow protobuf FreezeGraph
请注意,此方法仅在tensorflow_version <2

时有效
    # 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
                               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

下一步,您需要获取有关PB模型输出层的详细信息

from onnx_tf.backend import prepare

# load the model saved in onnx format
model_onnx = onnx.load('/content/mnist.onnx')

# prepare model for exporting to tensorFlow using tensorFlow backend
tf_rep = prepare(model_onnx)
print(tf_rep.run(dummy_input)) # run sample inference of your model
print(tf_rep.inputs) # Input nodes to the model
print('-----')
print(tf_rep.outputs) # Output nodes from the model
print('-----')
print(tf_rep.tensor_dict) # All nodes in the model

# export tensorFlow backend to tensorflow tf file
tf_rep.export_graph('mnist.pb')

它输出的输出通过在第一个函数中指定的输出层的名称找到节点,例如在我的模型中是

print(tf_rep.tensor_dict) # All nodes in the model

现在记住张量的运算符add_10

从protobuf FrozenGraph转换模型的最好方法是使用官方的tensorflow lite convert documentations现在使用此代码将该模型从.pb文件转换为tflite模型

'main_output': <tf.Tensor 'add_10:0' shape=(1, 2) dtype=float32  

要选择最适合模型优化的选项,请参阅此官方指南

https://www.tensorflow.org/lite/guide/get_started#4_optimize_your_model_optional

请查看我的代码

答案 1 :(得分:1)

<块引用>

现在您可以直接在手机上运行 PyTorch 模型。查看 PyTorch Mobile 的文档 here

<块引用>

此答案适用于 TensorFlow 版本 2 或更高版本,
对于 TensorFlow 版本 1,点击 here

将模型从 protobuf freezeGraph 转换为 TFlite 的最佳方法是使用官方的 TensorFlow lite 转换器 documentation

根据 TensorFlow Docs,TocoConverter 已被弃用

<块引用>

该类 (tf.compat.v1.lite.TocoConverter) 已被弃用。请改用 lite.TFLiteConverter。

从 PyTorch 转换为 ONNX 模型

# 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
)

所以在你的情况下: 现在使用 onnx-tf

将此模型导出到 TensorFlow protobuf FreezeGraph

从 ONNX 转换为 TensorFlow freezeGraph

要转换模型,请从以下命令安装 onnx-tf

git clone https://github.com/onnx/onnx-tensorflow.git && cd onnx-tensorflow
pip install -e .

现在将 .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_saved_model('mnist.pb')
# 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